Skip to content

Commit 170cdf5

Browse files
iamKunalGuptaiamKunalGupta
committed
feat: try testing flakiness
1 parent 4e4a293 commit 170cdf5

File tree

1 file changed

+77
-68
lines changed

1 file changed

+77
-68
lines changed

‎flow/connectors/postgres/rds_test.go‎

Lines changed: 77 additions & 68 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ package connpostgres
22

33
import (
44
"os"
5+
"strconv"
56
"testing"
67

78
"github.com/aws/smithy-go/ptr"
@@ -11,82 +12,90 @@ import (
1112
)
1213

1314
func TestAwsRDSIAMAuthConnectForPostgres(t *testing.T) {
14-
t.Setenv("AWS_ACCESS_KEY_ID", os.Getenv("FLOW_TESTS_RDS_IAM_AUTH_AWS_ACCESS_KEY_ID"))
15-
t.Setenv("AWS_SECRET_ACCESS_KEY", os.Getenv("FLOW_TESTS_RDS_IAM_AUTH_AWS_SECRET_ACCESS_KEY"))
16-
t.Setenv("AWS_SESSION_TOKEN", os.Getenv("FLOW_TESTS_RDS_IAM_AUTH_AWS_SESSION_TOKEN"))
17-
host := os.Getenv("FLOW_TESTS_RDS_IAM_AUTH_HOST_POSTGRES")
18-
username := os.Getenv("FLOW_TESTS_RDS_IAM_AUTH_USERNAME_POSTGRES")
19-
postgresConnector, err := NewPostgresConnector(t.Context(),
20-
nil,
21-
&protos.PostgresConfig{
22-
Host: host,
23-
Database: "postgres",
24-
User: username,
25-
Port: 5432,
26-
AuthType: protos.PostgresAuthType_POSTGRES_IAM_AUTH,
27-
RequireTls: true, // Assumed that AWS Root CA is installed
28-
AwsAuth: &protos.AwsAuthenticationConfig{
29-
AuthType: protos.AwsIAMAuthConfigType_IAM_AUTH_ASSUME_ROLE,
30-
AuthConfig: &protos.AwsAuthenticationConfig_Role{
31-
Role: &protos.AWSAuthAssumeRoleConfig{
32-
AssumeRoleArn: os.Getenv("FLOW_TESTS_RDS_IAM_AUTH_ASSUME_ROLE"),
33-
ChainedRoleArn: ptr.String(os.Getenv("FLOW_TESTS_RDS_IAM_AUTH_CHAINED_ROLE")),
15+
for attempt := 0; attempt < 30; attempt++ {
16+
t.Run(t.Name()+"-attempt-"+strconv.Itoa(attempt), func(t *testing.T) {
17+
t.Setenv("AWS_ACCESS_KEY_ID", os.Getenv("FLOW_TESTS_RDS_IAM_AUTH_AWS_ACCESS_KEY_ID"))
18+
t.Setenv("AWS_SECRET_ACCESS_KEY", os.Getenv("FLOW_TESTS_RDS_IAM_AUTH_AWS_SECRET_ACCESS_KEY"))
19+
t.Setenv("AWS_SESSION_TOKEN", os.Getenv("FLOW_TESTS_RDS_IAM_AUTH_AWS_SESSION_TOKEN"))
20+
host := os.Getenv("FLOW_TESTS_RDS_IAM_AUTH_HOST_POSTGRES")
21+
username := os.Getenv("FLOW_TESTS_RDS_IAM_AUTH_USERNAME_POSTGRES")
22+
postgresConnector, err := NewPostgresConnector(t.Context(),
23+
nil,
24+
&protos.PostgresConfig{
25+
Host: host,
26+
Database: "postgres",
27+
User: username,
28+
Port: 5432,
29+
AuthType: protos.PostgresAuthType_POSTGRES_IAM_AUTH,
30+
RequireTls: true, // Assumed that AWS Root CA is installed
31+
AwsAuth: &protos.AwsAuthenticationConfig{
32+
AuthType: protos.AwsIAMAuthConfigType_IAM_AUTH_ASSUME_ROLE,
33+
AuthConfig: &protos.AwsAuthenticationConfig_Role{
34+
Role: &protos.AWSAuthAssumeRoleConfig{
35+
AssumeRoleArn: os.Getenv("FLOW_TESTS_RDS_IAM_AUTH_ASSUME_ROLE"),
36+
ChainedRoleArn: ptr.String(os.Getenv("FLOW_TESTS_RDS_IAM_AUTH_CHAINED_ROLE")),
37+
},
38+
},
3439
},
35-
},
36-
},
40+
})
41+
require.NoError(t, err)
42+
defer postgresConnector.Close()
43+
rows, err := postgresConnector.Conn().Query(t.Context(), "SELECT 1")
44+
require.NoError(t, err)
45+
rowCount := 0
46+
for rows.Next() {
47+
rowCount++
48+
var val int
49+
require.NoError(t, rows.Scan(&val))
50+
require.Equal(t, 1, val)
51+
}
52+
require.Equal(t, 1, rowCount)
3753
})
38-
require.NoError(t, err)
39-
defer postgresConnector.Close()
40-
rows, err := postgresConnector.Conn().Query(t.Context(), "SELECT 1")
41-
require.NoError(t, err)
42-
rowCount := 0
43-
for rows.Next() {
44-
rowCount++
45-
var val int
46-
require.NoError(t, rows.Scan(&val))
47-
require.Equal(t, 1, val)
4854
}
49-
require.Equal(t, 1, rowCount)
5055
}
5156

5257
func TestAwsRDSIAMAuthConnectForPostgresViaProxy(t *testing.T) {
53-
t.Setenv("AWS_ACCESS_KEY_ID", os.Getenv("FLOW_TESTS_RDS_IAM_AUTH_AWS_ACCESS_KEY_ID"))
54-
t.Setenv("AWS_SECRET_ACCESS_KEY", os.Getenv("FLOW_TESTS_RDS_IAM_AUTH_AWS_SECRET_ACCESS_KEY"))
55-
t.Setenv("AWS_SESSION_TOKEN", os.Getenv("FLOW_TESTS_RDS_IAM_AUTH_AWS_SESSION_TOKEN"))
56-
rdsHost := os.Getenv("FLOW_TESTS_RDS_IAM_AUTH_HOST_POSTGRES")
57-
proxyHost := os.Getenv("FLOW_TESTS_RDS_IAM_AUTH_HOST_POSTGRES_PROXY")
58+
for attempt := 0; attempt < 30; attempt++ {
59+
t.Run(t.Name()+"-attempt-"+strconv.Itoa(attempt), func(t *testing.T) {
60+
t.Setenv("AWS_ACCESS_KEY_ID", os.Getenv("FLOW_TESTS_RDS_IAM_AUTH_AWS_ACCESS_KEY_ID"))
61+
t.Setenv("AWS_SECRET_ACCESS_KEY", os.Getenv("FLOW_TESTS_RDS_IAM_AUTH_AWS_SECRET_ACCESS_KEY"))
62+
t.Setenv("AWS_SESSION_TOKEN", os.Getenv("FLOW_TESTS_RDS_IAM_AUTH_AWS_SESSION_TOKEN"))
63+
rdsHost := os.Getenv("FLOW_TESTS_RDS_IAM_AUTH_HOST_POSTGRES")
64+
proxyHost := os.Getenv("FLOW_TESTS_RDS_IAM_AUTH_HOST_POSTGRES_PROXY")
5865

59-
username := os.Getenv("FLOW_TESTS_RDS_IAM_AUTH_USERNAME_POSTGRES")
60-
postgresConnector, err := NewPostgresConnector(t.Context(),
61-
nil,
62-
&protos.PostgresConfig{
63-
Host: proxyHost,
64-
Port: 5432,
65-
User: username,
66-
Database: "postgres",
67-
TlsHost: rdsHost,
68-
RequireTls: true, // Assumed that AWS Root CA is installed
69-
AuthType: protos.PostgresAuthType_POSTGRES_IAM_AUTH,
70-
AwsAuth: &protos.AwsAuthenticationConfig{
71-
AuthType: protos.AwsIAMAuthConfigType_IAM_AUTH_ASSUME_ROLE,
72-
AuthConfig: &protos.AwsAuthenticationConfig_Role{
73-
Role: &protos.AWSAuthAssumeRoleConfig{
74-
AssumeRoleArn: os.Getenv("FLOW_TESTS_RDS_IAM_AUTH_ASSUME_ROLE"),
75-
ChainedRoleArn: ptr.String(os.Getenv("FLOW_TESTS_RDS_IAM_AUTH_CHAINED_ROLE")),
66+
username := os.Getenv("FLOW_TESTS_RDS_IAM_AUTH_USERNAME_POSTGRES")
67+
postgresConnector, err := NewPostgresConnector(t.Context(),
68+
nil,
69+
&protos.PostgresConfig{
70+
Host: proxyHost,
71+
Port: 5432,
72+
User: username,
73+
Database: "postgres",
74+
TlsHost: rdsHost,
75+
RequireTls: true, // Assumed that AWS Root CA is installed
76+
AuthType: protos.PostgresAuthType_POSTGRES_IAM_AUTH,
77+
AwsAuth: &protos.AwsAuthenticationConfig{
78+
AuthType: protos.AwsIAMAuthConfigType_IAM_AUTH_ASSUME_ROLE,
79+
AuthConfig: &protos.AwsAuthenticationConfig_Role{
80+
Role: &protos.AWSAuthAssumeRoleConfig{
81+
AssumeRoleArn: os.Getenv("FLOW_TESTS_RDS_IAM_AUTH_ASSUME_ROLE"),
82+
ChainedRoleArn: ptr.String(os.Getenv("FLOW_TESTS_RDS_IAM_AUTH_CHAINED_ROLE")),
83+
},
84+
},
7685
},
77-
},
78-
},
86+
})
87+
require.NoError(t, err)
88+
defer postgresConnector.Close()
89+
rows, err := postgresConnector.Conn().Query(t.Context(), "SELECT 1")
90+
require.NoError(t, err)
91+
rowCount := 0
92+
for rows.Next() {
93+
rowCount++
94+
var val int
95+
require.NoError(t, rows.Scan(&val))
96+
require.Equal(t, 1, val)
97+
}
98+
require.Equal(t, 1, rowCount)
7999
})
80-
require.NoError(t, err)
81-
defer postgresConnector.Close()
82-
rows, err := postgresConnector.Conn().Query(t.Context(), "SELECT 1")
83-
require.NoError(t, err)
84-
rowCount := 0
85-
for rows.Next() {
86-
rowCount++
87-
var val int
88-
require.NoError(t, rows.Scan(&val))
89-
require.Equal(t, 1, val)
90100
}
91-
require.Equal(t, 1, rowCount)
92101
}

0 commit comments

Comments
 (0)