[release-15.0] Flaky tests: Fix wrangler tests (#13568)

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

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