Skip to content

Commit 42d5956

Browse files
kgala2kgala2
authored
test: update ip types for e2e tests (#2427)
1 parent f2dbaf8 commit 42d5956

File tree

5 files changed

+203
-97
lines changed

5 files changed

+203
-97
lines changed

‎tests/connection_test.go‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,9 @@ func removeAuthEnvVar(t *testing.T) (*oauth2.Token, string, func()) {
4343
if err != nil {
4444
t.Errorf("failed to get token: %v", err)
4545
}
46+
if *ipType != "public" {
47+
return tok, "", func() {}
48+
}
4649
path, ok := os.LookupEnv("GOOGLE_APPLICATION_CREDENTIALS")
4750
if !ok {
4851
t.Fatalf("GOOGLE_APPLICATION_CREDENTIALS was not set in the environment")

‎tests/fuse_test.go‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,9 @@ import (
2626
)
2727

2828
func TestPostgresFUSEConnect(t *testing.T) {
29+
if v := os.Getenv("IP_TYPE"); v == "private" || v == "psc" {
30+
t.Skipf("skipping test because IP_TYPE is set to %v", v)
31+
}
2932
if testing.Short() {
3033
t.Skip("skipping Postgres integration tests")
3134
}

‎tests/mysql_test.go‎

Lines changed: 79 additions & 31 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,29 @@ 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+
default:
73+
return args
74+
}
75+
}
76+
5977
func TestMySQLTCP(t *testing.T) {
6078
if testing.Short() {
6179
t.Skip("skipping MySQL integration tests")
6280
}
6381
requireMySQLVars(t)
64-
proxyConnTest(t, []string{*mysqlConnName}, "mysql", mysqlDSN())
82+
// Prepare the initial arguments
83+
args := []string{*mysqlConnName}
84+
// Add the IP type flag using the helper
85+
args = AddIPTypeFlag(args)
86+
// Run the test
87+
proxyConnTest(t, args, "mysql", mysqlDSN())
6588
}
6689

6790
func TestMySQLUnix(t *testing.T) {
@@ -82,8 +105,12 @@ func TestMySQLUnix(t *testing.T) {
82105
Addr: proxy.UnixAddress(tmpDir, *mysqlConnName),
83106
Net: "unix",
84107
}
85-
proxyConnTest(t,
86-
[]string{"--unix-socket", tmpDir, *mysqlConnName}, "mysql", cfg.FormatDSN())
108+
// Prepare the initial arguments
109+
args := []string{"--unix-socket", tmpDir, *mysqlConnName}
110+
// Add the IP type flag using the helper
111+
args = AddIPTypeFlag(args)
112+
// Run the test
113+
proxyConnTest(t, args, "mysql", cfg.FormatDSN())
87114
}
88115

89116
func TestMySQLImpersonation(t *testing.T) {
@@ -92,10 +119,15 @@ func TestMySQLImpersonation(t *testing.T) {
92119
}
93120
requireMySQLVars(t)
94121

95-
proxyConnTest(t, []string{
122+
// Prepare the initial arguments
123+
args := []string{
96124
"--impersonate-service-account", *impersonatedUser,
97-
*mysqlConnName},
98-
"mysql", mysqlDSN())
125+
*mysqlConnName,
126+
}
127+
// Add the IP type flag using the helper
128+
args = AddIPTypeFlag(args)
129+
// Run the test
130+
proxyConnTest(t, args, "mysql", mysqlDSN())
99131
}
100132

101133
func TestMySQLAuthentication(t *testing.T) {
@@ -104,7 +136,10 @@ func TestMySQLAuthentication(t *testing.T) {
104136
}
105137
requireMySQLVars(t)
106138

107-
creds := keyfile(t)
139+
var creds string
140+
if *ipType == "public" {
141+
creds = keyfile(t)
142+
}
108143
tok, path, cleanup := removeAuthEnvVar(t)
109144
defer cleanup()
110145

@@ -123,32 +158,42 @@ func TestMySQLAuthentication(t *testing.T) {
123158
"--impersonate-service-account", *impersonatedUser,
124159
*mysqlConnName},
125160
},
126-
{
127-
desc: "with credentials file",
128-
args: []string{"--credentials-file", path, *mysqlConnName},
129-
},
130-
{
131-
desc: "with credentials file and impersonation",
132-
args: []string{
133-
"--credentials-file", path,
134-
"--impersonate-service-account", *impersonatedUser,
135-
*mysqlConnName},
136-
},
137-
{
138-
desc: "with credentials JSON",
139-
args: []string{"--json-credentials", string(creds), *mysqlConnName},
140-
},
141-
{
142-
desc: "with credentials JSON and impersonation",
143-
args: []string{
144-
"--json-credentials", string(creds),
145-
"--impersonate-service-account", *impersonatedUser,
146-
*mysqlConnName},
147-
},
161+
}
162+
if *ipType == "public" {
163+
additionaTcs := []struct {
164+
desc string
165+
args []string
166+
}{
167+
{
168+
desc: "with credentials file",
169+
args: []string{"--credentials-file", path, *mysqlConnName},
170+
},
171+
{
172+
desc: "with credentials file and impersonation",
173+
args: []string{
174+
"--credentials-file", path,
175+
"--impersonate-service-account", *impersonatedUser,
176+
*mysqlConnName,
177+
},
178+
},
179+
{
180+
desc: "with credentials JSON",
181+
args: []string{"--json-credentials", string(creds), *mysqlConnName},
182+
},
183+
{
184+
desc: "with credentials JSON and impersonation",
185+
args: []string{
186+
"--json-credentials", string(creds),
187+
"--impersonate-service-account", *impersonatedUser,
188+
*mysqlConnName,
189+
},
190+
},
191+
}
192+
tcs = append(tcs, additionaTcs...)
148193
}
149194
for _, tc := range tcs {
150195
t.Run(tc.desc, func(t *testing.T) {
151-
proxyConnTest(t, tc.args, "mysql", mysqlDSN())
196+
proxyConnTest(t, AddIPTypeFlag(tc.args), "mysql", mysqlDSN())
152197
})
153198
}
154199
}
@@ -157,6 +202,9 @@ func TestMySQLGcloudAuth(t *testing.T) {
157202
if testing.Short() {
158203
t.Skip("skipping MySQL integration tests")
159204
}
205+
if v := os.Getenv("IP_TYPE"); v == "private" || v == "psc" {
206+
t.Skipf("skipping test because IP_TYPE is set to %v", v)
207+
}
160208
requireMySQLVars(t)
161209

162210
tcs := []struct {
@@ -177,7 +225,7 @@ func TestMySQLGcloudAuth(t *testing.T) {
177225
}
178226
for _, tc := range tcs {
179227
t.Run(tc.desc, func(t *testing.T) {
180-
proxyConnTest(t, tc.args, "mysql", mysqlDSN())
228+
proxyConnTest(t, AddIPTypeFlag(tc.args), "mysql", mysqlDSN())
181229
})
182230
}
183231
}

‎tests/postgres_test.go‎

Lines changed: 63 additions & 35 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) {
@@ -112,7 +124,10 @@ func TestPostgresAuthentication(t *testing.T) {
112124
}
113125
requirePostgresVars(t)
114126

115-
creds := keyfile(t)
127+
var creds string
128+
if *ipType == "public" {
129+
creds = keyfile(t)
130+
}
116131
tok, path, cleanup := removeAuthEnvVar(t)
117132
defer cleanup()
118133

@@ -131,32 +146,42 @@ func TestPostgresAuthentication(t *testing.T) {
131146
"--impersonate-service-account", *impersonatedUser,
132147
*postgresConnName},
133148
},
134-
{
135-
desc: "with credentials file",
136-
args: []string{"--credentials-file", path, *postgresConnName},
137-
},
138-
{
139-
desc: "with credentials file and impersonation",
140-
args: []string{
141-
"--credentials-file", path,
142-
"--impersonate-service-account", *impersonatedUser,
143-
*postgresConnName},
144-
},
145-
{
146-
desc: "with credentials JSON",
147-
args: []string{"--json-credentials", string(creds), *postgresConnName},
148-
},
149-
{
150-
desc: "with credentials JSON and impersonation",
151-
args: []string{
152-
"--json-credentials", string(creds),
153-
"--impersonate-service-account", *impersonatedUser,
154-
*postgresConnName},
155-
},
149+
}
150+
if *ipType == "public" {
151+
additionalTcs := []struct {
152+
desc string
153+
args []string
154+
}{
155+
{
156+
desc: "with credentials file",
157+
args: []string{"--credentials-file", path, *postgresConnName},
158+
},
159+
{
160+
desc: "with credentials file and impersonation",
161+
args: []string{
162+
"--credentials-file", path,
163+
"--impersonate-service-account", *impersonatedUser,
164+
*postgresConnName,
165+
},
166+
},
167+
{
168+
desc: "with credentials JSON",
169+
args: []string{"--json-credentials", string(creds), *postgresConnName},
170+
},
171+
{
172+
desc: "with credentials JSON and impersonation",
173+
args: []string{
174+
"--json-credentials", string(creds),
175+
"--impersonate-service-account", *impersonatedUser,
176+
*postgresConnName,
177+
},
178+
},
179+
}
180+
tcs = append(tcs, additionalTcs...)
156181
}
157182
for _, tc := range tcs {
158183
t.Run(tc.desc, func(t *testing.T) {
159-
proxyConnTest(t, tc.args, "pgx", postgresDSN())
184+
proxyConnTest(t, AddIPTypeFlag(tc.args), "pgx", postgresDSN())
160185
})
161186
}
162187
}
@@ -165,6 +190,9 @@ func TestPostgresGcloudAuth(t *testing.T) {
165190
if testing.Short() {
166191
t.Skip("skipping Postgres integration tests")
167192
}
193+
if v := os.Getenv("IP_TYPE"); v == "private" || v == "psc" {
194+
t.Skipf("skipping test because IP_TYPE is set to %v", v)
195+
}
168196
requirePostgresVars(t)
169197

170198
tcs := []struct {
@@ -185,7 +213,7 @@ func TestPostgresGcloudAuth(t *testing.T) {
185213
}
186214
for _, tc := range tcs {
187215
t.Run(tc.desc, func(t *testing.T) {
188-
proxyConnTest(t, tc.args, "pgx", postgresDSN())
216+
proxyConnTest(t, AddIPTypeFlag(tc.args), "pgx", postgresDSN())
189217
})
190218
}
191219

@@ -231,7 +259,7 @@ func TestPostgresIAMDBAuthn(t *testing.T) {
231259
}
232260
for _, tc := range tcs {
233261
t.Run(tc.desc, func(t *testing.T) {
234-
proxyConnTest(t, tc.args, "pgx", tc.dsn)
262+
proxyConnTest(t, AddIPTypeFlag(tc.args), "pgx", tc.dsn)
235263
})
236264
}
237265
}
@@ -272,7 +300,7 @@ func TestPostgresCustomerCAS(t *testing.T) {
272300
}
273301
for _, tc := range tcs {
274302
t.Run(tc.desc, func(t *testing.T) {
275-
proxyConnTest(t, tc.args, "pgx", tc.dsn)
303+
proxyConnTest(t, AddIPTypeFlag(tc.args), "pgx", tc.dsn)
276304
})
277305
}
278306
}

0 commit comments

Comments
 (0)