Skip to content

Commit 8a96fb1

Browse files
NirajNandreNirajNandre
authored
fix(tests/postgres): implement uuid-based isolation and reliable resource cleanup (#2377)
…atements ## Description This PR resolves integration test resource leaks and memory exhaustion by implementing per-test isolation and reliable teardown logic. **Files Updated:** - **tests/common.go**: Refactored CleanupPostgresTables to target specific uniqueID resources. - **tests/alloydbpg/alloydb_pg_integration_test.go**: Integrated t.Cleanup and unique ID logging. ## PR Checklist > Thank you for opening a Pull Request! Before submitting your PR, there are a > few things you can do to make sure it goes smoothly: - [ ] Make sure you reviewed [CONTRIBUTING.md](https://github.com/googleapis/genai-toolbox/blob/main/CONTRIBUTING.md) - [ ] Make sure to open an issue as a [bug/issue](https://github.com/googleapis/genai-toolbox/issues/new/choose) before writing your code! That way we can discuss the change, evaluate designs, and agree on the general idea - [ ] Ensure the tests and linter pass - [ ] Code coverage does not decrease (if any source code was changed) - [ ] Appropriate docs were updated (if necessary) - [ ] Make sure to add `!` if this involve a breaking change 🛠️ Fixes #<2463>
1 parent d4ff6be commit 8a96fb1

File tree

5 files changed

+59
-33
lines changed

5 files changed

+59
-33
lines changed

‎tests/alloydbpg/alloydb_pg_integration_test.go‎

Lines changed: 12 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -129,13 +129,18 @@ func TestAlloyDBPgToolEndpoints(t *testing.T) {
129129
t.Fatalf("unable to create AlloyDB connection pool: %s", err)
130130
}
131131

132-
// cleanup test environment
133-
tests.CleanupPostgresTables(t, ctx, pool)
134-
135-
// create table name with UUID
136-
tableNameParam := "param_table_" + strings.ReplaceAll(uuid.New().String(), "-", "")
137-
tableNameAuth := "auth_table_" + strings.ReplaceAll(uuid.New().String(), "-", "")
138-
tableNameTemplateParam := "template_param_table_" + strings.ReplaceAll(uuid.New().String(), "-", "")
132+
// Generate a unique ID
133+
uniqueID := strings.ReplaceAll(uuid.New().String(), "-", "")
134+
135+
// This will execute after all tool tests complete (success, fail, or t.Fatal)
136+
t.Cleanup(func() {
137+
tests.CleanupPostgresTables(t, context.Background(), pool, uniqueID)
138+
})
139+
140+
//Create table names using the UUID
141+
tableNameParam := "param_table_" + uniqueID
142+
tableNameAuth := "auth_table_" + uniqueID
143+
tableNameTemplateParam := "template_param_table_" + uniqueID
139144

140145
// set up data for param tool
141146
createParamTableStmt, insertParamTableStmt, paramToolStmt, idParamToolStmt, nameParamToolStmt, arrayToolStmt, paramTestParams := tests.GetPostgresSQLParamToolInfo(tableNameParam)

‎tests/cloudsqlpg/cloud_sql_pg_integration_test.go‎

Lines changed: 11 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -114,13 +114,18 @@ func TestCloudSQLPgSimpleToolEndpoints(t *testing.T) {
114114
t.Fatalf("unable to create Cloud SQL connection pool: %s", err)
115115
}
116116

117-
// cleanup test environment
118-
tests.CleanupPostgresTables(t, ctx, pool)
117+
// Generate a unique ID
118+
uniqueID := strings.ReplaceAll(uuid.New().String(), "-", "")
119119

120-
// create table name with UUID
121-
tableNameParam := "param_table_" + strings.ReplaceAll(uuid.New().String(), "-", "")
122-
tableNameAuth := "auth_table_" + strings.ReplaceAll(uuid.New().String(), "-", "")
123-
tableNameTemplateParam := "template_param_table_" + strings.ReplaceAll(uuid.New().String(), "-", "")
120+
// This will execute after all tool tests complete (success, fail, or t.Fatal)
121+
t.Cleanup(func() {
122+
tests.CleanupPostgresTables(t, context.Background(), pool, uniqueID)
123+
})
124+
125+
//Create table names using the UUID
126+
tableNameParam := "param_table_" + uniqueID
127+
tableNameAuth := "auth_table_" + uniqueID
128+
tableNameTemplateParam := "template_param_table_" + uniqueID
124129

125130
// set up data for param tool
126131
createParamTableStmt, insertParamTableStmt, paramToolStmt, idParamToolStmt, nameParamToolStmt, arrayToolStmt, paramTestParams := tests.GetPostgresSQLParamToolInfo(tableNameParam)

‎tests/cockroachdb/cockroachdb_integration_test.go‎

Lines changed: 11 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -109,13 +109,18 @@ func TestCockroachDB(t *testing.T) {
109109
}
110110
t.Logf("✅ Connected to: %s", version)
111111

112-
// cleanup test environment
113-
tests.CleanupPostgresTables(t, ctx, pool)
112+
// Generate a unique ID
113+
uniqueID := strings.ReplaceAll(uuid.New().String(), "-", "")
114114

115-
// create table names with UUID suffix
116-
tableNameParam := "param_table_" + strings.ReplaceAll(uuid.New().String(), "-", "")
117-
tableNameAuth := "auth_table_" + strings.ReplaceAll(uuid.New().String(), "-", "")
118-
tableNameTemplateParam := "template_param_table_" + strings.ReplaceAll(uuid.New().String(), "-", "")
115+
// This will execute after all tool tests complete (success, fail, or t.Fatal)
116+
t.Cleanup(func() {
117+
tests.CleanupPostgresTables(t, context.Background(), pool, uniqueID)
118+
})
119+
120+
//Create table names using the UUID
121+
tableNameParam := "param_table_" + uniqueID
122+
tableNameAuth := "auth_table_" + uniqueID
123+
tableNameTemplateParam := "template_param_table_" + uniqueID
119124

120125
// set up data for param tool (using CockroachDB explicit INT primary keys)
121126
createParamTableStmt, insertParamTableStmt, paramToolStmt, idParamToolStmt, nameParamToolStmt, arrayToolStmt, paramTestParams := tests.GetCockroachDBParamToolInfo(tableNameParam)

‎tests/common.go‎

Lines changed: 13 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -972,14 +972,15 @@ func TestCloudSQLMySQL_IPTypeParsingFromYAML(t *testing.T) {
972972
}
973973

974974
// Finds and drops all tables in a postgres database.
975-
func CleanupPostgresTables(t *testing.T, ctx context.Context, pool *pgxpool.Pool) {
976-
query := `
977-
SELECT table_name FROM information_schema.tables
978-
WHERE table_schema = 'public' AND table_type = 'BASE TABLE';`
975+
func CleanupPostgresTables(t *testing.T, ctx context.Context, pool *pgxpool.Pool, uniqueID string) {
979976

980-
rows, err := pool.Query(ctx, query)
977+
query := `
978+
SELECT table_name FROM information_schema.tables
979+
WHERE table_schema = 'public'
980+
AND table_name LIKE $1;`
981+
rows, err := pool.Query(ctx, query, "%\\_"+uniqueID)
981982
if err != nil {
982-
t.Fatalf("Failed to query for all tables in 'public' schema: %v", err)
983+
t.Fatalf("Failed to query for tables for uniqueID %s in 'public' schema: %v", uniqueID, err)
983984
}
984985
defer rows.Close()
985986

@@ -994,13 +995,18 @@ func CleanupPostgresTables(t *testing.T, ctx context.Context, pool *pgxpool.Pool
994995
}
995996

996997
if len(tablesToDrop) == 0 {
998+
t.Logf("INTEGRATION CLEANUP: No tables found for uniqueID: %s", uniqueID)
997999
return
9981000
}
9991001

1002+
t.Logf("INTEGRATION CLEANUP: Dropping %d isolated tables: %v", len(tablesToDrop), tablesToDrop)
1003+
10001004
dropQuery := fmt.Sprintf("DROP TABLE IF EXISTS %s CASCADE;", strings.Join(tablesToDrop, ", "))
10011005

10021006
if _, err := pool.Exec(ctx, dropQuery); err != nil {
1003-
t.Fatalf("Failed to drop all tables in 'public' schema: %v", err)
1007+
t.Fatalf("Failed to drop tables for uniqueID %s in 'public' schema: %v", uniqueID, err)
1008+
} else {
1009+
t.Logf("INTEGRATION CLEANUP SUCCESS: Wiped all tables for uniqueID: %s", uniqueID)
10041010
}
10051011
}
10061012

‎tests/postgres/postgres_integration_test.go‎

Lines changed: 12 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -93,13 +93,18 @@ func TestPostgres(t *testing.T) {
9393
t.Fatalf("unable to create postgres connection pool: %s", err)
9494
}
9595

96-
// cleanup test environment
97-
tests.CleanupPostgresTables(t, ctx, pool)
98-
99-
// create table name with UUID
100-
tableNameParam := "param_table_" + strings.ReplaceAll(uuid.New().String(), "-", "")
101-
tableNameAuth := "auth_table_" + strings.ReplaceAll(uuid.New().String(), "-", "")
102-
tableNameTemplateParam := "template_param_table_" + strings.ReplaceAll(uuid.New().String(), "-", "")
96+
// Generate a unique ID
97+
uniqueID := strings.ReplaceAll(uuid.New().String(), "-", "")
98+
99+
// This will execute after all tool tests complete (success, fail, or t.Fatal)
100+
t.Cleanup(func() {
101+
tests.CleanupPostgresTables(t, context.Background(), pool, uniqueID)
102+
})
103+
104+
//Create table names using the UUID
105+
tableNameParam := "param_table_" + uniqueID
106+
tableNameAuth := "auth_table_" + uniqueID
107+
tableNameTemplateParam := "template_param_table_" + uniqueID
103108

104109
// set up data for param tool
105110
createParamTableStmt, insertParamTableStmt, paramToolStmt, idParamToolStmt, nameParamToolStmt, arrayToolStmt, paramTestParams := tests.GetPostgresSQLParamToolInfo(tableNameParam)

0 commit comments

Comments
 (0)