Flaky tests: Fix wrangler tests (#13568)

Signed-off-by: Manan Gupta <[email protected]>

go/vt/wrangler/testlib/copy_schema_shard_test.go

@@ -155,6 +155,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)
 	}

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)
 

go/vt/wrangler/testlib/planned_reparent_shard_test.go

@@ -878,26 +878,6 @@ func TestPlannedReparentShardPromoteReplicaFail(t *testing.T) {
 	assert.True(t, oldPrimary.FakeMysqlDaemon.ReadOnly, "oldPrimary.FakeMysqlDaemon.ReadOnly")
 }
 
-// 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)
-		}
-	}
-}
-
 // TestPlannedReparentShardSamePrimary tests PRS with oldPrimary works correctly
 // Simulate failure of previous PRS and oldPrimary is ReadOnly
 // Verify that primary correctly gets set to ReadWrite

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)
+		}
+	}
+}