From 07221ad651dd8cd06437601e759722939dab6556 Mon Sep 17 00:00:00 2001
From: "vitess-bot[bot]" <108069721+vitess-bot[bot]@users.noreply.github.com>
Date: Thu, 20 Jul 2023 17:43:46 +0530
Subject: [PATCH] Flaky tests: Fix wrangler tests (#13568)

Signed-off-by: Manan Gupta <manan@planetscale.com>
---
 .../testlib/copy_schema_shard_test.go         |  4 ++
 .../testlib/external_reparent_test.go         |  3 +
 go/vt/wrangler/testlib/utils.go               | 57 +++++++++++++++++++
 3 files changed, 64 insertions(+)
 create mode 100644 go/vt/wrangler/testlib/utils.go

diff --git a/go/vt/wrangler/testlib/copy_schema_shard_test.go b/go/vt/wrangler/testlib/copy_schema_shard_test.go
index 0f62db641b6..d192aa3d42e 100644
--- a/go/vt/wrangler/testlib/copy_schema_shard_test.go
+++ b/go/vt/wrangler/testlib/copy_schema_shard_test.go
@@ -171,6 +171,10 @@ func copySchema(t *testing.T, useShardAsSource bool) {
 	if useShardAsSource {
 		source = "ks/-80"
 	}
+
+	// PrimaryAlias in the shard record is updated asynchronously, so we should wait for it to succeed.
+	waitForShardPrimary(t, wr, destinationPrimary.Tablet)
+
 	if err := vp.Run([]string{"CopySchemaShard", "--include-views", source, "ks/-40"}); err != nil {
 		t.Fatalf("CopySchemaShard failed: %v", err)
 	}
diff --git a/go/vt/wrangler/testlib/external_reparent_test.go b/go/vt/wrangler/testlib/external_reparent_test.go
index 00fd10f703f..8d928926550 100644
--- a/go/vt/wrangler/testlib/external_reparent_test.go
+++ b/go/vt/wrangler/testlib/external_reparent_test.go
@@ -550,6 +550,9 @@ func TestRPCTabletExternallyReparentedDemotesPrimaryToConfiguredTabletType(t *te
 		}
 	}
 
+	// PrimaryAlias in the shard record is updated asynchronously, so we should wait for it to succeed.
+	waitForShardPrimary(t, wr, newPrimary.Tablet)
+
 	shardInfo, err := ts.GetShard(context.Background(), newPrimary.Tablet.Keyspace, newPrimary.Tablet.Shard)
 	assert.NoError(t, err)
 
diff --git a/go/vt/wrangler/testlib/utils.go b/go/vt/wrangler/testlib/utils.go
new file mode 100644
index 00000000000..ae32e5036bc
--- /dev/null
+++ b/go/vt/wrangler/testlib/utils.go
@@ -0,0 +1,57 @@
+package testlib
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+
+	topodatapb "vitess.io/vitess/go/vt/proto/topodata"
+	"vitess.io/vitess/go/vt/topo/topoproto"
+	"vitess.io/vitess/go/vt/wrangler"
+)
+
+// waitForTabletType waits for the given tablet type to be reached.
+func waitForTabletType(t *testing.T, wr *wrangler.Wrangler, tabletAlias *topodatapb.TabletAlias, tabletType topodatapb.TabletType) {
+	timeout := time.After(15 * time.Second)
+	for {
+		tablet, err := wr.TopoServer().GetTablet(context.Background(), tabletAlias)
+		require.NoError(t, err)
+		if tablet.Type == tabletType {
+			return
+		}
+
+		select {
+		case <-timeout:
+			t.Fatalf("%s didn't reach the tablet type %v", topoproto.TabletAliasString(tabletAlias), tabletType.String())
+			return
+		default:
+			time.Sleep(100 * time.Millisecond)
+		}
+	}
+}
+
+// waitForShardPrimary waits for the shard record to be upto date such that it has the given primary.
+func waitForShardPrimary(t *testing.T, wr *wrangler.Wrangler, primaryTablet *topodatapb.Tablet) {
+	timeout := time.After(15 * time.Second)
+	for {
+		si, err := wr.TopoServer().GetShard(context.Background(), primaryTablet.Keyspace, primaryTablet.Shard)
+		require.NoError(t, err)
+		if topoproto.TabletAliasEqual(si.PrimaryAlias, primaryTablet.Alias) {
+			return
+		}
+
+		select {
+		case <-timeout:
+			t.Fatalf("%s/%s didn't see the tablet %v become the primary, instead it is %v",
+				primaryTablet.Keyspace, primaryTablet.Shard,
+				topoproto.TabletAliasString(primaryTablet.Alias),
+				topoproto.TabletAliasString(si.PrimaryAlias),
+			)
+			return
+		default:
+			time.Sleep(100 * time.Millisecond)
+		}
+	}
+}