feat: use a singe timeout for all the deletions

Signed-off-by: Manan Gupta <[email protected]>
vitessio · Nov 15, 2023 · 9cd716f · 9cd716f
1 parent 64124f7
commit 9cd716f
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 4 deletions.
diff --git a/go/test/endtoend/utils/utils.go b/go/test/endtoend/utils/utils.go
@@ -283,11 +283,10 @@ func WaitForKsError(t *testing.T, vtgateProcess cluster.VtgateProcess, ks string
 }
 
 // WaitForTableDeletions waits for a table to be deleted
-func WaitForTableDeletions(t *testing.T, vtgateProcess cluster.VtgateProcess, ks, tbl string) error {
-	timeout := time.After(60 * time.Second)
+func WaitForTableDeletions(ctx context.Context, t *testing.T, vtgateProcess cluster.VtgateProcess, ks, tbl string) error {
 	for {
 		select {
-		case <-timeout:
+		case <-ctx.Done():
 			return fmt.Errorf("schema tracking still found the table '%s'", tbl)
 		default:
 			res, err := vtgateProcess.ReadVSchema()

diff --git a/go/test/endtoend/vtgate/foreignkey/stress/fk_stress_test.go b/go/test/endtoend/vtgate/foreignkey/stress/fk_stress_test.go
@@ -705,8 +705,10 @@ func createInitialSchema(t *testing.T, tcase *testCase) {
 		}
 	})
 	t.Run("waiting for vschema deletions to apply", func(t *testing.T) {
+		timeoutCtx, cancel := context.WithTimeout(ctx, 1*time.Minute)
+		defer cancel()
 		for _, tableName := range tableNames {
-			err := utils.WaitForTableDeletions(t, clusterInstance.VtgateProcess, keyspaceName, tableName)
+			err := utils.WaitForTableDeletions(timeoutCtx, t, clusterInstance.VtgateProcess, keyspaceName, tableName)
 			require.NoError(t, err)
 		}
 	})