From 10a7424167ce25866ee8acca5ce192d9c7b5367f Mon Sep 17 00:00:00 2001
From: Harshit Gangal <harshit@planetscale.com>
Date: Mon, 20 Nov 2023 10:29:56 +0530
Subject: [PATCH] addressed review comments

Signed-off-by: Harshit Gangal <harshit@planetscale.com>
---
 go/test/endtoend/utils/utils.go               | 39 ++++++++++---------
 go/test/endtoend/vtgate/foreignkey/fk_test.go | 17 +-------
 .../foreignkey/stress/fk_stress_test.go       |  3 +-
 3 files changed, 24 insertions(+), 35 deletions(-)

diff --git a/go/test/endtoend/utils/utils.go b/go/test/endtoend/utils/utils.go
index c837278c7a6..9cf9c767fc7 100644
--- a/go/test/endtoend/utils/utils.go
+++ b/go/test/endtoend/utils/utils.go
@@ -255,6 +255,21 @@ func WaitForAuthoritative(t *testing.T, ks, tbl string, readVSchema func() (*int
 }
 
 // WaitForKsError waits for the ks error field to be populated and returns it.
+func WaitForKsError(t *testing.T, vtgateProcess cluster.VtgateProcess, ks string) string {
+	var errString string
+	WaitForVschemaCondition(t, vtgateProcess, ks, func(t *testing.T, keyspace map[string]interface{}) bool {
+		ksErr, fieldPresent := keyspace["error"]
+		if !fieldPresent {
+			return false
+		}
+		var ok bool
+		errString, ok = ksErr.(string)
+		return ok
+	})
+	return errString
+}
+
+// WaitForVschemaCondition waits for the condition to be true
 func WaitForVschemaCondition(t *testing.T, vtgateProcess cluster.VtgateProcess, ks string, conditionMet func(t *testing.T, keyspace map[string]interface{}) bool) {
 	timeout := time.After(60 * time.Second)
 	for {
@@ -275,24 +290,12 @@ func WaitForVschemaCondition(t *testing.T, vtgateProcess cluster.VtgateProcess,
 }
 
 // WaitForTableDeletions waits for a table to be deleted
-func WaitForTableDeletions(ctx context.Context, t *testing.T, vtgateProcess cluster.VtgateProcess, ks, tbl string) error {
-	for {
-		select {
-		case <-ctx.Done():
-			return fmt.Errorf("schema tracking still found the table '%s'", tbl)
-		default:
-			res, err := vtgateProcess.ReadVSchema()
-			require.NoError(t, err, res)
-			keyspacesMap := convertToMap(*res)["keyspaces"]
-			ksMap := convertToMap(keyspacesMap)[ks]
-			tablesMap := convertToMap(ksMap)["tables"]
-			_, isPresent := convertToMap(tablesMap)[tbl]
-			if !isPresent {
-				return nil
-			}
-			time.Sleep(100 * time.Millisecond)
-		}
-	}
+func WaitForTableDeletions(ctx context.Context, t *testing.T, vtgateProcess cluster.VtgateProcess, ks, tbl string) {
+	WaitForVschemaCondition(t, vtgateProcess, ks, func(t *testing.T, keyspace map[string]interface{}) bool {
+		tablesMap := keyspace["tables"]
+		_, isPresent := convertToMap(tablesMap)[tbl]
+		return !isPresent
+	})
 }
 
 // WaitForColumn waits for a table's column to be present
diff --git a/go/test/endtoend/vtgate/foreignkey/fk_test.go b/go/test/endtoend/vtgate/foreignkey/fk_test.go
index 2043833276b..6d17976c815 100644
--- a/go/test/endtoend/vtgate/foreignkey/fk_test.go
+++ b/go/test/endtoend/vtgate/foreignkey/fk_test.go
@@ -979,22 +979,9 @@ func TestCyclicFks(t *testing.T) {
 	// Create a cyclic foreign key constraint.
 	utils.Exec(t, mcmp.VtConn, "alter table fk_t10 add constraint test_cyclic_fks foreign key (col) references fk_t12 (col) on delete cascade on update cascade")
 
-	matchKsError := func(t *testing.T, keyspace map[string]interface{}) bool {
-		ksErr, fieldPresent := keyspace["error"]
-		if !fieldPresent {
-			return false
-		}
-		errString, isErr := ksErr.(string)
-		if !isErr {
-			return false
-		}
-		// Make sure Vschema has the error for cyclic foreign keys.
-		assert.Contains(t, errString, "VT09019: keyspace 'uks' has cyclic foreign keys")
-		return true
-	}
-
 	// Wait for schema-tracking to be complete.
-	utils.WaitForVschemaCondition(t, clusterInstance.VtgateProcess, unshardedKs, matchKsError)
+	errString := utils.WaitForKsError(t, clusterInstance.VtgateProcess, unshardedKs)
+	assert.Contains(t, errString, "VT09019: keyspace 'uks' has cyclic foreign keys")
 
 	// Ensure that the Vitess database is originally empty
 	ensureDatabaseState(t, mcmp.VtConn, true)
diff --git a/go/test/endtoend/vtgate/foreignkey/stress/fk_stress_test.go b/go/test/endtoend/vtgate/foreignkey/stress/fk_stress_test.go
index 00b7d9ca0cc..e31d2a22e04 100644
--- a/go/test/endtoend/vtgate/foreignkey/stress/fk_stress_test.go
+++ b/go/test/endtoend/vtgate/foreignkey/stress/fk_stress_test.go
@@ -707,8 +707,7 @@ func createInitialSchema(t *testing.T, tcase *testCase) {
 		timeoutCtx, cancel := context.WithTimeout(ctx, 1*time.Minute)
 		defer cancel()
 		for _, tableName := range tableNames {
-			err := utils.WaitForTableDeletions(timeoutCtx, t, clusterInstance.VtgateProcess, keyspaceName, tableName)
-			require.NoError(t, err)
+			utils.WaitForTableDeletions(timeoutCtx, t, clusterInstance.VtgateProcess, keyspaceName, tableName)
 		}
 	})
 	t.Run("creating tables", func(t *testing.T) {