Skip to content

Commit e3bfd70

Browse files
kgala2kgala2
committed
chore: update ip type for e2e tests
1 parent b139a27 commit e3bfd70

3 files changed

Lines changed: 100 additions & 28 deletions

File tree

‎tests/mysql_test.go‎

Lines changed: 48 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,12 @@ var (
2929
mysqlUser = flag.String("mysql_user", os.Getenv("MYSQL_USER"), "Name of database user.")
3030
mysqlPass = flag.String("mysql_pass", os.Getenv("MYSQL_PASS"), "Password for the database user; be careful when entering a password on the command line (it may go into your terminal's history).")
3131
mysqlDB = flag.String("mysql_db", os.Getenv("MYSQL_DB"), "Name of the database to connect to.")
32+
ipType = flag.String("ip_type", func() string {
33+
if v := os.Getenv("IP_TYPE"); v != "" {
34+
return v
35+
}
36+
return "public"
37+
}(), "IP type of the instance to connect to, can be public, private or psc")
3238
)
3339

3440
func requireMySQLVars(t *testing.T) {
@@ -56,12 +62,33 @@ func mysqlDSN() string {
5662
return cfg.FormatDSN()
5763
}
5864

65+
// addIPTypeFlag appends the correct flag based on the ipType variable.
66+
func addIPTypeFlag(args []string) []string {
67+
switch *ipType {
68+
case "private":
69+
return append(args, "--private-ip")
70+
case "psc":
71+
return append(args, "--psc")
72+
// "public" is the default and doesn't require a flag
73+
case "public":
74+
return args
75+
default:
76+
// Or handle unknown ipType values as needed, maybe log a warning?
77+
return args
78+
}
79+
}
80+
5981
func TestMySQLTCP(t *testing.T) {
6082
if testing.Short() {
6183
t.Skip("skipping MySQL integration tests")
6284
}
6385
requireMySQLVars(t)
64-
proxyConnTest(t, []string{*mysqlConnName}, "mysql", mysqlDSN())
86+
// Prepare the initial arguments
87+
args := []string{*mysqlConnName}
88+
// Add the IP type flag using the helper
89+
args = addIPTypeFlag(args)
90+
// Run the test
91+
proxyConnTest(t, args, "mysql", mysqlDSN())
6592
}
6693

6794
func TestMySQLUnix(t *testing.T) {
@@ -82,8 +109,12 @@ func TestMySQLUnix(t *testing.T) {
82109
Addr: proxy.UnixAddress(tmpDir, *mysqlConnName),
83110
Net: "unix",
84111
}
85-
proxyConnTest(t,
86-
[]string{"--unix-socket", tmpDir, *mysqlConnName}, "mysql", cfg.FormatDSN())
112+
// Prepare the initial arguments
113+
args := []string{"--unix-socket", tmpDir, *mysqlConnName}
114+
// Add the IP type flag using the helper
115+
args = addIPTypeFlag(args)
116+
// Run the test
117+
proxyConnTest(t, args, "mysql", cfg.FormatDSN())
87118
}
88119

89120
func TestMySQLImpersonation(t *testing.T) {
@@ -92,10 +123,15 @@ func TestMySQLImpersonation(t *testing.T) {
92123
}
93124
requireMySQLVars(t)
94125

95-
proxyConnTest(t, []string{
126+
// Prepare the initial arguments
127+
args := []string{
96128
"--impersonate-service-account", *impersonatedUser,
97-
*mysqlConnName},
98-
"mysql", mysqlDSN())
129+
*mysqlConnName,
130+
}
131+
// Add the IP type flag using the helper
132+
args = addIPTypeFlag(args)
133+
// Run the test
134+
proxyConnTest(t, args, "mysql", mysqlDSN())
99135
}
100136

101137
func TestMySQLAuthentication(t *testing.T) {
@@ -148,7 +184,9 @@ func TestMySQLAuthentication(t *testing.T) {
148184
}
149185
for _, tc := range tcs {
150186
t.Run(tc.desc, func(t *testing.T) {
151-
proxyConnTest(t, tc.args, "mysql", mysqlDSN())
187+
// Add the IP type flag using the helper
188+
argsWithIPType := addIPTypeFlag(tc.args)
189+
proxyConnTest(t, argsWithIPType, "mysql", mysqlDSN())
152190
})
153191
}
154192
}
@@ -177,7 +215,9 @@ func TestMySQLGcloudAuth(t *testing.T) {
177215
}
178216
for _, tc := range tcs {
179217
t.Run(tc.desc, func(t *testing.T) {
180-
proxyConnTest(t, tc.args, "mysql", mysqlDSN())
218+
// Add the IP type flag using the helper
219+
argsWithIPType := addIPTypeFlag(tc.args)
220+
proxyConnTest(t, argsWithIPType, "mysql", mysqlDSN())
181221
})
182222
}
183223
}

‎tests/postgres_test.go‎

Lines changed: 32 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -60,8 +60,12 @@ func TestPostgresTCP(t *testing.T) {
6060
t.Skip("skipping Postgres integration tests")
6161
}
6262
requirePostgresVars(t)
63-
64-
proxyConnTest(t, []string{*postgresConnName}, "pgx", postgresDSN())
63+
// Prepare the initial arguments
64+
args := []string{*postgresConnName}
65+
// Add the IP type flag using the helper
66+
args = addIPTypeFlag(args)
67+
// Run the test
68+
proxyConnTest(t, args, "pgx", postgresDSN())
6569
}
6670

6771
func TestPostgresUnix(t *testing.T) {
@@ -78,8 +82,12 @@ func TestPostgresUnix(t *testing.T) {
7882
proxy.UnixAddress(tmpDir, *postgresConnName),
7983
*postgresUser, *postgresPass, *postgresDB)
8084

81-
proxyConnTest(t,
82-
[]string{"--unix-socket", tmpDir, *postgresConnName}, "pgx", dsn)
85+
// Prepare the initial arguments
86+
args := []string{"--unix-socket", tmpDir, *postgresConnName}
87+
// Add the IP type flag using the helper
88+
args = addIPTypeFlag(args)
89+
// Run the test
90+
proxyConnTest(t, args, "pgx", dsn)
8391
}
8492

8593
func createTempDir(t *testing.T) (string, func()) {
@@ -99,11 +107,15 @@ func TestPostgresImpersonation(t *testing.T) {
99107
t.Skip("skipping Postgres integration tests")
100108
}
101109
requirePostgresVars(t)
102-
103-
proxyConnTest(t, []string{
110+
// Prepare the initial arguments
111+
args := []string{
104112
"--impersonate-service-account", *impersonatedUser,
105-
*postgresConnName},
106-
"pgx", postgresDSN())
113+
*postgresConnName,
114+
}
115+
// Add the IP type flag using the helper
116+
args = addIPTypeFlag(args)
117+
// Run the test
118+
proxyConnTest(t, args, "pgx", postgresDSN())
107119
}
108120

109121
func TestPostgresAuthentication(t *testing.T) {
@@ -156,7 +168,9 @@ func TestPostgresAuthentication(t *testing.T) {
156168
}
157169
for _, tc := range tcs {
158170
t.Run(tc.desc, func(t *testing.T) {
159-
proxyConnTest(t, tc.args, "pgx", postgresDSN())
171+
// Add the IP type flag using the helper
172+
argsWithIPType := addIPTypeFlag(tc.args)
173+
proxyConnTest(t, argsWithIPType, "pgx", postgresDSN())
160174
})
161175
}
162176
}
@@ -185,7 +199,9 @@ func TestPostgresGcloudAuth(t *testing.T) {
185199
}
186200
for _, tc := range tcs {
187201
t.Run(tc.desc, func(t *testing.T) {
188-
proxyConnTest(t, tc.args, "pgx", postgresDSN())
202+
// Add the IP type flag using the helper
203+
argsWithIPType := addIPTypeFlag(tc.args)
204+
proxyConnTest(t, argsWithIPType, "pgx", postgresDSN())
189205
})
190206
}
191207

@@ -231,7 +247,9 @@ func TestPostgresIAMDBAuthn(t *testing.T) {
231247
}
232248
for _, tc := range tcs {
233249
t.Run(tc.desc, func(t *testing.T) {
234-
proxyConnTest(t, tc.args, "pgx", tc.dsn)
250+
// Add the IP type flag using the helper
251+
argsWithIPType := addIPTypeFlag(tc.args)
252+
proxyConnTest(t, argsWithIPType, "pgx", tc.dsn)
235253
})
236254
}
237255
}
@@ -272,7 +290,9 @@ func TestPostgresCustomerCAS(t *testing.T) {
272290
}
273291
for _, tc := range tcs {
274292
t.Run(tc.desc, func(t *testing.T) {
275-
proxyConnTest(t, tc.args, "pgx", tc.dsn)
293+
// Add the IP type flag using the helper
294+
argsWithIPType := addIPTypeFlag(tc.args)
295+
proxyConnTest(t, argsWithIPType, "pgx", tc.dsn)
276296
})
277297
}
278298
}

‎tests/sqlserver_test.go‎

Lines changed: 20 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -54,20 +54,28 @@ func TestSQLServerTCP(t *testing.T) {
5454
t.Skip("skipping SQL Server integration tests")
5555
}
5656
requireSQLServerVars(t)
57-
58-
proxyConnTest(t, []string{*sqlserverConnName}, "sqlserver", sqlserverDSN())
57+
// Prepare the initial arguments
58+
args := []string{*sqlserverConnName}
59+
// Add the IP type flag using the helper
60+
args = addIPTypeFlag(args)
61+
// Run the test
62+
proxyConnTest(t, args, "sqlserver", sqlserverDSN())
5963
}
6064

6165
func TestSQLServerImpersonation(t *testing.T) {
6266
if testing.Short() {
6367
t.Skip("skipping SQL Server integration tests")
6468
}
6569
requireSQLServerVars(t)
66-
67-
proxyConnTest(t, []string{
70+
// Prepare the initial arguments
71+
args := []string{
6872
"--impersonate-service-account", *impersonatedUser,
69-
*sqlserverConnName},
70-
"sqlserver", sqlserverDSN())
73+
*sqlserverConnName,
74+
}
75+
// Add the IP type flag using the helper
76+
args = addIPTypeFlag(args)
77+
// Run the test
78+
proxyConnTest(t, args, "sqlserver", sqlserverDSN())
7179
}
7280

7381
func TestSQLServerAuthentication(t *testing.T) {
@@ -120,7 +128,9 @@ func TestSQLServerAuthentication(t *testing.T) {
120128
}
121129
for _, tc := range tcs {
122130
t.Run(tc.desc, func(t *testing.T) {
123-
proxyConnTest(t, tc.args, "sqlserver", sqlserverDSN())
131+
// Add the IP type flag using the helper
132+
argsWithIPType := addIPTypeFlag(tc.args)
133+
proxyConnTest(t, argsWithIPType, "sqlserver", sqlserverDSN())
124134
})
125135
}
126136
}
@@ -149,7 +159,9 @@ func TestSQLServerGcloudAuth(t *testing.T) {
149159
}
150160
for _, tc := range tcs {
151161
t.Run(tc.desc, func(t *testing.T) {
152-
proxyConnTest(t, tc.args, "sqlserver", sqlserverDSN())
162+
// Add the IP type flag using the helper
163+
argsWithIPType := addIPTypeFlag(tc.args)
164+
proxyConnTest(t, argsWithIPType, "sqlserver", sqlserverDSN())
153165
})
154166
}
155167
}

0 commit comments

Comments
 (0)