From eda15e6a377d9fc55e8277241b1268714022c82d Mon Sep 17 00:00:00 2001
From: Rohit Nayak <57520317+rohit-nayak-ps@users.noreply.github.com>
Date: Thu, 7 Dec 2023 19:42:04 +0100
Subject: [PATCH] Flaky TestFKExtWorkflow: fix Foreign Key stress test
 flakiness (#14714)

---
 go/test/endtoend/vreplication/fk_ext_test.go | 34 ++++++++++++--------
 1 file changed, 21 insertions(+), 13 deletions(-)

diff --git a/go/test/endtoend/vreplication/fk_ext_test.go b/go/test/endtoend/vreplication/fk_ext_test.go
index 665f8052486..a06fafb257e 100644
--- a/go/test/endtoend/vreplication/fk_ext_test.go
+++ b/go/test/endtoend/vreplication/fk_ext_test.go
@@ -194,17 +194,8 @@ func TestFKExt(t *testing.T) {
 
 }
 
-// compareRowCounts compares the row counts for the parent and child tables in the source and target shards. In addition to vdiffs,
-// it is another check to ensure that both tables have the same number of rows in the source and target shards after load generation
-// has stopped.
-func compareRowCounts(t *testing.T, keyspace string, sourceShards, targetShards []string) error {
-	log.Infof("Comparing row counts for keyspace %s, source shards: %v, target shards: %v", keyspace, sourceShards, targetShards)
-	lg.Stop()
-	defer lg.Start()
-	if err := waitForCondition("load generator to stop", func() bool { return lg.State() == LoadGeneratorStateStopped }, 10*time.Second); err != nil {
-		return err
-	}
-
+// checkRowCounts checks that the parent and child tables in the source and target shards have the same number of rows.
+func checkRowCounts(t *testing.T, keyspace string, sourceShards, targetShards []string) bool {
 	sourceTabs := make(map[string]*cluster.VttabletProcess)
 	targetTabs := make(map[string]*cluster.VttabletProcess)
 	for _, shard := range sourceShards {
@@ -239,9 +230,26 @@ func compareRowCounts(t *testing.T, keyspace string, sourceShards, targetShards
 	log.Infof("Source parent count: %d, child count: %d, target parent count: %d, child count: %d.",
 		sourceParentCount, sourceChildCount, targetParentCount, targetChildCount)
 	if sourceParentCount != targetParentCount || sourceChildCount != targetChildCount {
-		return fmt.Errorf(fmt.Sprintf("source and target row counts do not match; source parent count: %d, target parent count: %d, source child count: %d, target child count: %d",
-			sourceParentCount, targetParentCount, sourceChildCount, targetChildCount))
+		log.Infof("Row counts do not match for keyspace %s, source shards: %v, target shards: %v", keyspace, sourceShards, targetShards)
+		return false
+	}
+	return true
+}
+
+// compareRowCounts compares the row counts for the parent and child tables in the source and target shards. In addition to vdiffs,
+// it is another check to ensure that both tables have the same number of rows in the source and target shards after load generation
+// has stopped.
+func compareRowCounts(t *testing.T, keyspace string, sourceShards, targetShards []string) error {
+	log.Infof("Comparing row counts for keyspace %s, source shards: %v, target shards: %v", keyspace, sourceShards, targetShards)
+	lg.Stop()
+	defer lg.Start()
+	if err := waitForCondition("load generator to stop", func() bool { return lg.State() == LoadGeneratorStateStopped }, 10*time.Second); err != nil {
+		return err
 	}
+	if err := waitForCondition("matching row counts", func() bool { return checkRowCounts(t, keyspace, sourceShards, targetShards) }, 30*time.Second); err != nil {
+		return err
+	}
+
 	return nil
 }