From 135a6a8fb3a46bccd3ae39b3a0930956041307e1 Mon Sep 17 00:00:00 2001
From: Shlomi Noach <2607934+shlomi-noach@users.noreply.github.com>
Date: Mon, 17 Jun 2024 13:29:21 +0300
Subject: [PATCH] CI: wait-for rather than 'assume' in Online DDL flow (#16210)

Signed-off-by: Shlomi Noach <2607934+shlomi-noach@users.noreply.github.com>
---
 .../onlineddl/flow/onlineddl_flow_test.go     | 25 ++++++++++++++-----
 1 file changed, 19 insertions(+), 6 deletions(-)

diff --git a/go/test/endtoend/onlineddl/flow/onlineddl_flow_test.go b/go/test/endtoend/onlineddl/flow/onlineddl_flow_test.go
index 772a5fa6fd0..aed3efcde2f 100644
--- a/go/test/endtoend/onlineddl/flow/onlineddl_flow_test.go
+++ b/go/test/endtoend/onlineddl/flow/onlineddl_flow_test.go
@@ -304,13 +304,26 @@ func TestSchemaChange(t *testing.T) {
 				}
 				waitForThrottleCheckStatus(t, throttlerapp.OnlineDDLName, primaryTablet, http.StatusOK)
 			})
-			t.Run("additional wait", func(t *testing.T) {
-				// Waiting just so that we generate more DMLs, and give migration/vreplication
+			t.Run("apply more DML", func(t *testing.T) {
+				// Looking to run a substantial amount of DML, giving vreplication
 				// more "opportunities" to throttle or to make progress.
-				select {
-				case <-time.After(3 * time.Second):
-				case <-ctx.Done():
-					require.Fail(t, "context cancelled")
+				ctx, cancel := context.WithTimeout(context.Background(), time.Minute)
+				defer cancel()
+				ticker := time.NewTicker(time.Second)
+				defer ticker.Stop()
+
+				startDML := totalAppliedDML.Load()
+				for {
+					appliedDML := totalAppliedDML.Load()
+					if appliedDML-startDML >= int64(maxTableRows) {
+						// We have generated enough DMLs
+						return
+					}
+					select {
+					case <-ticker.C:
+					case <-ctx.Done():
+						require.Fail(t, "timeout waiting for applied DML")
+					}
 				}
 			})
 			t.Run("validate applied DML", func(t *testing.T) {