From 8a96fb1a8874baa3688e566f3dea8a0912fcf2df Mon Sep 17 00:00:00 2001 From: Niraj Nandre Date: Tue, 17 Feb 2026 19:09:53 +0530 Subject: [PATCH] fix(tests/postgres): implement uuid-based isolation and reliable resource cleanup (#2377) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit …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> --- .../alloydbpg/alloydb_pg_integration_test.go | 17 ++++++++++------ .../cloud_sql_pg_integration_test.go | 17 ++++++++++------ .../cockroachdb_integration_test.go | 17 ++++++++++------ tests/common.go | 20 ++++++++++++------- tests/postgres/postgres_integration_test.go | 17 ++++++++++------ 5 files changed, 57 insertions(+), 31 deletions(-) diff --git a/tests/alloydbpg/alloydb_pg_integration_test.go b/tests/alloydbpg/alloydb_pg_integration_test.go index 92542567907..dc0751665ba 100644 --- a/tests/alloydbpg/alloydb_pg_integration_test.go +++ b/tests/alloydbpg/alloydb_pg_integration_test.go @@ -129,13 +129,18 @@ func TestAlloyDBPgToolEndpoints(t *testing.T) { t.Fatalf("unable to create AlloyDB connection pool: %s", err) } - // cleanup test environment - tests.CleanupPostgresTables(t, ctx, pool) + // Generate a unique ID + uniqueID := strings.ReplaceAll(uuid.New().String(), "-", "") - // create table name with UUID - tableNameParam := "param_table_" + strings.ReplaceAll(uuid.New().String(), "-", "") - tableNameAuth := "auth_table_" + strings.ReplaceAll(uuid.New().String(), "-", "") - tableNameTemplateParam := "template_param_table_" + strings.ReplaceAll(uuid.New().String(), "-", "") + // This will execute after all tool tests complete (success, fail, or t.Fatal) + t.Cleanup(func() { + tests.CleanupPostgresTables(t, context.Background(), pool, uniqueID) + }) + + //Create table names using the UUID + tableNameParam := "param_table_" + uniqueID + tableNameAuth := "auth_table_" + uniqueID + tableNameTemplateParam := "template_param_table_" + uniqueID // set up data for param tool createParamTableStmt, insertParamTableStmt, paramToolStmt, idParamToolStmt, nameParamToolStmt, arrayToolStmt, paramTestParams := tests.GetPostgresSQLParamToolInfo(tableNameParam) diff --git a/tests/cloudsqlpg/cloud_sql_pg_integration_test.go b/tests/cloudsqlpg/cloud_sql_pg_integration_test.go index e6340cb8717..0408d5e371a 100644 --- a/tests/cloudsqlpg/cloud_sql_pg_integration_test.go +++ b/tests/cloudsqlpg/cloud_sql_pg_integration_test.go @@ -114,13 +114,18 @@ func TestCloudSQLPgSimpleToolEndpoints(t *testing.T) { t.Fatalf("unable to create Cloud SQL connection pool: %s", err) } - // cleanup test environment - tests.CleanupPostgresTables(t, ctx, pool) + // Generate a unique ID + uniqueID := strings.ReplaceAll(uuid.New().String(), "-", "") - // create table name with UUID - tableNameParam := "param_table_" + strings.ReplaceAll(uuid.New().String(), "-", "") - tableNameAuth := "auth_table_" + strings.ReplaceAll(uuid.New().String(), "-", "") - tableNameTemplateParam := "template_param_table_" + strings.ReplaceAll(uuid.New().String(), "-", "") + // This will execute after all tool tests complete (success, fail, or t.Fatal) + t.Cleanup(func() { + tests.CleanupPostgresTables(t, context.Background(), pool, uniqueID) + }) + + //Create table names using the UUID + tableNameParam := "param_table_" + uniqueID + tableNameAuth := "auth_table_" + uniqueID + tableNameTemplateParam := "template_param_table_" + uniqueID // set up data for param tool createParamTableStmt, insertParamTableStmt, paramToolStmt, idParamToolStmt, nameParamToolStmt, arrayToolStmt, paramTestParams := tests.GetPostgresSQLParamToolInfo(tableNameParam) diff --git a/tests/cockroachdb/cockroachdb_integration_test.go b/tests/cockroachdb/cockroachdb_integration_test.go index 3fff3f7716f..464b82238e4 100644 --- a/tests/cockroachdb/cockroachdb_integration_test.go +++ b/tests/cockroachdb/cockroachdb_integration_test.go @@ -109,13 +109,18 @@ func TestCockroachDB(t *testing.T) { } t.Logf("✅ Connected to: %s", version) - // cleanup test environment - tests.CleanupPostgresTables(t, ctx, pool) + // Generate a unique ID + uniqueID := strings.ReplaceAll(uuid.New().String(), "-", "") - // create table names with UUID suffix - tableNameParam := "param_table_" + strings.ReplaceAll(uuid.New().String(), "-", "") - tableNameAuth := "auth_table_" + strings.ReplaceAll(uuid.New().String(), "-", "") - tableNameTemplateParam := "template_param_table_" + strings.ReplaceAll(uuid.New().String(), "-", "") + // This will execute after all tool tests complete (success, fail, or t.Fatal) + t.Cleanup(func() { + tests.CleanupPostgresTables(t, context.Background(), pool, uniqueID) + }) + + //Create table names using the UUID + tableNameParam := "param_table_" + uniqueID + tableNameAuth := "auth_table_" + uniqueID + tableNameTemplateParam := "template_param_table_" + uniqueID // set up data for param tool (using CockroachDB explicit INT primary keys) createParamTableStmt, insertParamTableStmt, paramToolStmt, idParamToolStmt, nameParamToolStmt, arrayToolStmt, paramTestParams := tests.GetCockroachDBParamToolInfo(tableNameParam) diff --git a/tests/common.go b/tests/common.go index 480143dbce6..db36144eb37 100644 --- a/tests/common.go +++ b/tests/common.go @@ -972,14 +972,15 @@ func TestCloudSQLMySQL_IPTypeParsingFromYAML(t *testing.T) { } // Finds and drops all tables in a postgres database. -func CleanupPostgresTables(t *testing.T, ctx context.Context, pool *pgxpool.Pool) { - query := ` - SELECT table_name FROM information_schema.tables - WHERE table_schema = 'public' AND table_type = 'BASE TABLE';` +func CleanupPostgresTables(t *testing.T, ctx context.Context, pool *pgxpool.Pool, uniqueID string) { - rows, err := pool.Query(ctx, query) + query := ` + SELECT table_name FROM information_schema.tables + WHERE table_schema = 'public' + AND table_name LIKE $1;` + rows, err := pool.Query(ctx, query, "%\\_"+uniqueID) if err != nil { - t.Fatalf("Failed to query for all tables in 'public' schema: %v", err) + t.Fatalf("Failed to query for tables for uniqueID %s in 'public' schema: %v", uniqueID, err) } defer rows.Close() @@ -994,13 +995,18 @@ func CleanupPostgresTables(t *testing.T, ctx context.Context, pool *pgxpool.Pool } if len(tablesToDrop) == 0 { + t.Logf("INTEGRATION CLEANUP: No tables found for uniqueID: %s", uniqueID) return } + t.Logf("INTEGRATION CLEANUP: Dropping %d isolated tables: %v", len(tablesToDrop), tablesToDrop) + dropQuery := fmt.Sprintf("DROP TABLE IF EXISTS %s CASCADE;", strings.Join(tablesToDrop, ", ")) if _, err := pool.Exec(ctx, dropQuery); err != nil { - t.Fatalf("Failed to drop all tables in 'public' schema: %v", err) + t.Fatalf("Failed to drop tables for uniqueID %s in 'public' schema: %v", uniqueID, err) + } else { + t.Logf("INTEGRATION CLEANUP SUCCESS: Wiped all tables for uniqueID: %s", uniqueID) } } diff --git a/tests/postgres/postgres_integration_test.go b/tests/postgres/postgres_integration_test.go index 06cb19310d0..174f4287e34 100644 --- a/tests/postgres/postgres_integration_test.go +++ b/tests/postgres/postgres_integration_test.go @@ -93,13 +93,18 @@ func TestPostgres(t *testing.T) { t.Fatalf("unable to create postgres connection pool: %s", err) } - // cleanup test environment - tests.CleanupPostgresTables(t, ctx, pool) + // Generate a unique ID + uniqueID := strings.ReplaceAll(uuid.New().String(), "-", "") - // create table name with UUID - tableNameParam := "param_table_" + strings.ReplaceAll(uuid.New().String(), "-", "") - tableNameAuth := "auth_table_" + strings.ReplaceAll(uuid.New().String(), "-", "") - tableNameTemplateParam := "template_param_table_" + strings.ReplaceAll(uuid.New().String(), "-", "") + // This will execute after all tool tests complete (success, fail, or t.Fatal) + t.Cleanup(func() { + tests.CleanupPostgresTables(t, context.Background(), pool, uniqueID) + }) + + //Create table names using the UUID + tableNameParam := "param_table_" + uniqueID + tableNameAuth := "auth_table_" + uniqueID + tableNameTemplateParam := "template_param_table_" + uniqueID // set up data for param tool createParamTableStmt, insertParamTableStmt, paramToolStmt, idParamToolStmt, nameParamToolStmt, arrayToolStmt, paramTestParams := tests.GetPostgresSQLParamToolInfo(tableNameParam)