style-refactor waitForRowCountInTablet()

Signed-off-by: Shlomi Noach <2607934+shlomi-noach@users.noreply.github.com>

go/test/endtoend/vreplication/helper_test.go

@@ -246,25 +246,35 @@ func waitForRowCount(t *testing.T, conn *mysql.Conn, database string, table stri
 	}
 }
 
-func waitForRowCountInTablet(t *testing.T, vttablet *cluster.VttabletProcess, database string, table string, want int) {
-	query := fmt.Sprintf("select count(*) from %s", table)
-	wantRes := fmt.Sprintf("[[INT64(%d)]]", want)
+func waitForRowCountInTablet(t *testing.T, vttablet *cluster.VttabletProcess, database string, table string, want int64) {
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTimeout)
+	defer cancel()
+	ticker := time.NewTicker(defaultTick)
+	defer ticker.Stop()
+
+	query := fmt.Sprintf("select count(*) as c from %s", table)
 	timer := time.NewTimer(defaultTimeout)
 	defer timer.Stop()
 	for {
 		qr, err := vttablet.QueryTablet(query, database, true)
 		require.NoError(t, err)
 		require.NotNil(t, qr)
-		if wantRes == fmt.Sprintf("%v", qr.Rows) {
+		row := qr.Named().Row()
+		require.NotNil(t, row)
+		got := row.AsInt64("c", 0)
+		require.LessOrEqual(t, got, want)
+		if got == want {
 			log.Infof("waitForRowCountInTablet: found %d rows in table %s on tablet %s", want, table, vttablet.Name)
 			return
 		}
 		select {
-		case <-timer.C:
-			require.FailNow(t, fmt.Sprintf("table %q did not reach the expected number of rows (%d) on tablet %q before the timeout of %s; last seen result: %v",
-				table, want, vttablet.Name, defaultTimeout, qr.Rows))
-		default:
-			time.Sleep(defaultTick)
+		case <-ctx.Done():
+			require.FailNow(
+				t, fmt.Sprintf("table %q did not reach the expected number of rows (%d) on tablet %q before the timeout of %s; last seen result: %v",
+					table, want, vttablet.Name, defaultTimeout, qr.Rows),
+			)
+			return
+		case <-ticker.C:
 		}
 	}
 }

go/test/endtoend/vreplication/multi_tenant_test.go

@@ -249,7 +249,7 @@ func TestMultiTenantSimple(t *testing.T) {
 	sourceTablet := vc.getPrimaryTablet(t, sourceKeyspace, "0")
 	require.NotNil(t, sourceTablet)
 	// Wait for the rows to be reverse replicated to the source keyspace.
-	waitForRowCountInTablet(t, sourceTablet, sourceKeyspace, "t1", int(lastIndex))
+	waitForRowCountInTablet(t, sourceTablet, sourceKeyspace, "t1", lastIndex)
 
 	mt.Complete()
 	require.Zero(t, len(getKeyspaceRoutingRules(t, vc).Rules))

go/test/endtoend/vreplication/vreplication_test.go

@@ -1209,7 +1209,7 @@ func reshard(t *testing.T, ksName string, tableName string, workflow string, sou
 			if tablets[tabletName] == nil {
 				continue
 			}
-			waitForRowCountInTablet(t, tablets[tabletName], ksName, tableName, count)
+			waitForRowCountInTablet(t, tablets[tabletName], ksName, tableName, int64(count))
 		}
 	})
 }