diff --git a/tests/alloydbpg/alloydb_pg_integration_test.go b/tests/alloydbpg/alloydb_pg_integration_test.go index 9254256790..dc0751665b 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 e6340cb871..0408d5e371 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 3fff3f7716..464b82238e 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 480143dbce..db36144eb3 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 06cb19310d..174f4287e3 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)