Skip to content

Commit a769f15

Browse files
NirajNandreNirajNandre
authored
fix(tests/Bigtable): implement uuid-based isolation and reliable resource cleanup (#2880)
## Description - **UUID-based Isolation**: Refactored the Bigtable integration tests to generate a single `uniqueID` at the start of the test run, which is used as a suffix for all temporary table names. - **Centralized Cleanup Utility**: Added the `CleanupBigtableTables` function to `tests/common.go` to identify and delete all tables matching the specific `uniqueID`. - **Reliable Teardown**: Replaced `defer` calls with the `t.Cleanup` hook and `context.Background()` to ensure that test-specific tables are consistently removed, even if the test fails or times out. **Why it was required**: - **Prevent Resource leaks**: Previous test runs frequently leaked tables when teardown logic was interrupted, leading to significant clutter in the toolbox-testing-438616 project. - **Resource Limit Near-Exhaustion**: The Bigtable instance was found to have **876** tables currently in use, approaching the standard cloud limit of **1,000 tables**. This cleanup is critical to maintain the health of the shared testing environment. - **Enable Concurrency**: Strict isolation ensures that parallel builds in CI/CD do not interfere with or accidentally delete resources belonging to other active test runs. Reference Doc - [https://docs.google.com/document/d/1LzOVRhuwU6z88FFE5ocyJ_7MIhY0kdpRFnuqU_yYAPg/edit?resourcekey=0-dcyFbd8jFgPgULqmeUlRjQ&tab=t.1rq6wvhsf9yn](https://docs.google.com/document/d/1LzOVRhuwU6z88FFE5ocyJ_7MIhY0kdpRFnuqU_yYAPg/edit?resourcekey=0-dcyFbd8jFgPgULqmeUlRjQ&tab=t.1rq6wvhsf9yn ) > Should include a concise description of the changes (bug or feature), it's > impact, along with a summary of the solution ## 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 #<issue_number_goes_here>
1 parent ddf409d commit a769f15

2 files changed

Lines changed: 59 additions & 35 deletions

File tree

‎tests/bigtable/bigtable_integration_test.go‎

Lines changed: 37 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -64,14 +64,33 @@ type TestRow struct {
6464

6565
func TestBigtableToolEndpoints(t *testing.T) {
6666
sourceConfig := getBigtableVars(t)
67-
ctx, cancel := context.WithTimeout(context.Background(), time.Minute)
68-
defer cancel()
67+
68+
uniqueID := strings.ReplaceAll(uuid.New().String(), "-", "")
69+
t.Logf("Starting Bigtable test with uniqueID: %s", uniqueID)
6970

7071
var args []string
7172

72-
tableName := "param_table" + strings.ReplaceAll(uuid.New().String(), "-", "")
73-
tableNameAuth := "auth_table_" + strings.ReplaceAll(uuid.New().String(), "-", "")
74-
tableNameTemplateParam := "tmpl_param_table_" + strings.ReplaceAll(uuid.New().String(), "-", "")
73+
// Initialize AdminClient to create or delete tables
74+
adminClient, err := bigtable.NewAdminClient(context.Background(), sourceConfig["project"].(string), sourceConfig["instance"].(string))
75+
if err != nil {
76+
t.Fatalf("Failed to create AdminClient: %v", err)
77+
}
78+
79+
t.Cleanup(func() {
80+
adminClient.Close()
81+
})
82+
83+
t.Cleanup(func() {
84+
t.Logf("Running global cleanup for uniqueID: %s", uniqueID)
85+
tests.CleanupBigtableTables(t, context.Background(), adminClient, uniqueID)
86+
})
87+
88+
ctx, cancel := context.WithTimeout(context.Background(), 7*time.Minute)
89+
defer cancel()
90+
91+
tableName := "param_table_" + uniqueID
92+
tableNameAuth := "auth_table_" + uniqueID
93+
tableNameTemplateParam := "tmpl_param_table_" + uniqueID
7594

7695
columnFamilyName := "cf"
7796
muts, rowKeys := getTestData(columnFamilyName)
@@ -85,18 +104,15 @@ func TestBigtableToolEndpoints(t *testing.T) {
85104
"SELECT TO_INT64(cf['id']) AS id, CAST(cf['name'] AS string) AS name FROM %s WHERE TO_INT64(cf['id']) IN UNNEST(@idArray) AND CAST(cf['name'] AS string) IN UNNEST(@nameArray);",
86105
tableName,
87106
)
88-
teardownTable1 := setupBtTable(t, ctx, sourceConfig["project"].(string), sourceConfig["instance"].(string), tableName, columnFamilyName, muts, rowKeys)
89-
defer teardownTable1(t)
107+
setupBtTable(t, adminClient, ctx, sourceConfig["project"].(string), sourceConfig["instance"].(string), tableName, columnFamilyName, muts, rowKeys)
90108

91109
// Do not change the shape of statement without checking tests/common_test.go.
92110
// The structure and value of seed data has to match https://github.com/googleapis/genai-toolbox/blob/4dba0df12dc438eca3cb476ef52aa17cdf232c12/tests/common_test.go#L200-L251
93111
authToolStatement := fmt.Sprintf("SELECT CAST(cf['name'] AS string) as name FROM %s WHERE CAST(cf['email'] AS string) = @email;", tableNameAuth)
94-
teardownTable2 := setupBtTable(t, ctx, sourceConfig["project"].(string), sourceConfig["instance"].(string), tableNameAuth, columnFamilyName, muts, rowKeys)
95-
defer teardownTable2(t)
112+
setupBtTable(t, adminClient, ctx, sourceConfig["project"].(string), sourceConfig["instance"].(string), tableNameAuth, columnFamilyName, muts, rowKeys)
96113

97114
mutsTmpl, rowKeysTmpl := getTestDataTemplateParam(columnFamilyName)
98-
teardownTableTmpl := setupBtTable(t, ctx, sourceConfig["project"].(string), sourceConfig["instance"].(string), tableNameTemplateParam, columnFamilyName, mutsTmpl, rowKeysTmpl)
99-
defer teardownTableTmpl(t)
115+
setupBtTable(t, adminClient, ctx, sourceConfig["project"].(string), sourceConfig["instance"].(string), tableNameTemplateParam, columnFamilyName, mutsTmpl, rowKeysTmpl)
100116

101117
// Write config into a file and pass it to command
102118
toolsFile := tests.GetToolsConfig(sourceConfig, BigtableToolType, paramTestStatement, idParamTestStatement, nameParamTestStatement, arrayTestStatement, authToolStatement)
@@ -225,63 +241,49 @@ func getTestDataTemplateParam(columnFamilyName string) ([]*bigtable.Mutation, []
225241
return muts, rowKeys
226242
}
227243

228-
func setupBtTable(t *testing.T, ctx context.Context, projectId string, instance string, tableName string, columnFamilyName string, muts []*bigtable.Mutation, rowKeys []string) func(*testing.T) {
229-
// Creating clients
230-
adminClient, err := bigtable.NewAdminClient(ctx, projectId, instance)
231-
if err != nil {
232-
t.Fatalf("NewAdminClient: %v", err)
233-
}
244+
func setupBtTable(t *testing.T, adminClient *bigtable.AdminClient, ctx context.Context, projectId string, instance string, tableName string, columnFamilyName string, muts []*bigtable.Mutation, rowKeys []string) {
234245

235246
client, err := bigtable.NewClient(ctx, projectId, instance)
236247
if err != nil {
237-
log.Fatalf("Could not create data operations client: %v", err)
248+
t.Fatalf("Could not create data operations client: %v", err)
238249
}
239250
defer client.Close()
240251

241252
// Creating tables
242253
tables, err := adminClient.Tables(ctx)
243254
if err != nil {
244-
log.Fatalf("Could not fetch table list: %v", err)
255+
t.Fatalf("Could not fetch table list: %v", err)
245256
}
246257

247258
if !slices.Contains(tables, tableName) {
248-
log.Printf("Creating table %s", tableName)
259+
t.Logf("Creating table %s", tableName)
249260
if err := adminClient.CreateTable(ctx, tableName); err != nil {
250-
log.Fatalf("Could not create table %s: %v", tableName, err)
261+
t.Fatalf("Could not create table %s: %v", tableName, err)
251262
}
252263
}
253264

254265
tblInfo, err := adminClient.TableInfo(ctx, tableName)
255266
if err != nil {
256-
log.Fatalf("Could not read info for table %s: %v", tableName, err)
267+
t.Fatalf("Could not read info for table %s: %v", tableName, err)
257268
}
258269

259270
// Creating column family
260271
if !slices.Contains(tblInfo.Families, columnFamilyName) {
261272
if err := adminClient.CreateColumnFamily(ctx, tableName, columnFamilyName); err != nil {
262-
log.Fatalf("Could not create column family %s: %v", columnFamilyName, err)
273+
t.Fatalf("Could not create column family %s: %v", columnFamilyName, err)
263274
}
264275
}
265276

266277
tbl := client.Open(tableName)
267278
rowErrs, err := tbl.ApplyBulk(ctx, rowKeys, muts)
268279
if err != nil {
269-
log.Fatalf("Could not apply bulk row mutation: %v", err)
280+
t.Fatalf("Could not apply bulk row mutation: %v", err)
270281
}
271282
if rowErrs != nil {
272283
for _, rowErr := range rowErrs {
273-
log.Printf("Error writing row: %v", rowErr)
274-
}
275-
log.Fatalf("Could not write some rows")
276-
}
277-
278-
// Writing data
279-
return func(t *testing.T) {
280-
// tear down test
281-
if err = adminClient.DeleteTable(ctx, tableName); err != nil {
282-
log.Fatalf("Teardown failed. Could not delete table %s: %v", tableName, err)
284+
t.Logf("Error writing row: %v", rowErr)
283285
}
284-
defer adminClient.Close()
286+
t.Fatalf("Could not write some rows")
285287
}
286288
}
287289

‎tests/common.go‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@ import (
2525
"testing"
2626

2727
"cloud.google.com/go/bigquery"
28+
"cloud.google.com/go/bigtable"
2829
"github.com/google/go-cmp/cmp"
2930
"github.com/googleapis/genai-toolbox/internal/server"
3031
"github.com/googleapis/genai-toolbox/internal/sources/cloudsqlmysql"
@@ -1102,3 +1103,24 @@ func CleanupBigQueryDatasets(t *testing.T, ctx context.Context, client *bigquery
11021103
}
11031104
}
11041105
}
1106+
1107+
// finds and deletes all tables in a Bigtable instance that match the uniqueID.
1108+
func CleanupBigtableTables(t *testing.T, ctx context.Context, adminClient *bigtable.AdminClient, uniqueID string) {
1109+
tables, err := adminClient.Tables(ctx)
1110+
if err != nil {
1111+
t.Errorf("INTEGRATION CLEANUP: Failed to list tables: %v", err)
1112+
return
1113+
}
1114+
1115+
for _, table := range tables {
1116+
// Delete tables containing our uniqueID
1117+
if strings.Contains(table, uniqueID) {
1118+
t.Logf("INTEGRATION CLEANUP: Deleting table %s", table)
1119+
if err := adminClient.DeleteTable(ctx, table); err != nil {
1120+
t.Errorf("INTEGRATION CLEANUP: Failed to delete table %s: %v", table, err)
1121+
} else {
1122+
t.Logf("INTEGRATION CLEANUP SUCCESS: Wiped table %s", table)
1123+
}
1124+
}
1125+
}
1126+
}

0 commit comments

Comments
 (0)