[vtctld/wrangler] Extract EmergencyReparentShard logic to dedicated struct and add unit tests

go/vt/vtadmin/api_test.go

@@ -727,7 +727,7 @@ func TestGetSchemas(t *testing.T) {
 
 					for _, tablet := range cts {
 						// AddTablet also adds the keyspace + shard for us.
-						testutil.AddTablet(context.Background(), t, toposerver, tablet.Tablet)
+						testutil.AddTablet(context.Background(), t, toposerver, tablet.Tablet, nil)
 
 						// Adds each SchemaDefinition to the fake TabletManagerClient, or nil
 						// if there are no schemas for that tablet. (All tablet aliases must

go/vt/vtctl/grpcvtctldserver/server_test.go

@@ -209,9 +209,7 @@ func TestChangeTabletType(t *testing.T) {
 			vtctld := NewVtctldServer(ts)
 			testutil.TabletManagerClient.Topo = ts
 
-			for _, tablet := range tt.tablets {
-				testutil.AddTablet(ctx, t, ts, tablet)
-			}
+			testutil.AddTablets(ctx, t, ts, nil, tt.tablets...)
 
 			resp, err := vtctld.ChangeTabletType(ctx, tt.req)
 			if tt.shouldErr {
@@ -257,7 +255,7 @@ func TestChangeTabletType(t *testing.T) {
 				Uid:  100,
 			},
 			Type: topodatapb.TabletType_REPLICA,
-		})
+		}, nil)
 
 		_, err := vtctld.ChangeTabletType(ctx, &vtctldatapb.ChangeTabletTypeRequest{
 			TabletAlias: &topodatapb.TabletAlias{
@@ -1364,7 +1362,7 @@ func TestDeleteShards(t *testing.T) {
 			vtctld := NewVtctldServer(ts)
 
 			testutil.AddShards(ctx, t, ts, tt.shards...)
-			testutil.AddTablets(ctx, t, ts, tt.tablets...)
+			testutil.AddTablets(ctx, t, ts, nil, tt.tablets...)
 			testutil.SetupReplicationGraphs(ctx, t, ts, tt.replicationGraphs...)
 			testutil.UpdateSrvKeyspaces(ctx, t, ts, tt.srvKeyspaces)
 
@@ -1853,9 +1851,7 @@ func TestDeleteTablets(t *testing.T) {
 			vtctld := NewVtctldServer(ts)
 
 			// Setup tablets and shards
-			for _, tablet := range tt.tablets {
-				testutil.AddTablet(ctx, t, ts, tablet)
-			}
+			testutil.AddTablets(ctx, t, ts, nil, tt.tablets...)
 
 			for key, updateFn := range tt.shardFieldUpdates {
 				ks, shard, err := topoproto.ParseKeyspaceShard(key)
@@ -2166,7 +2162,7 @@ func TestGetTablet(t *testing.T) {
 		Type:     topodatapb.TabletType_REPLICA,
 	}
 
-	testutil.AddTablet(ctx, t, ts, tablet)
+	testutil.AddTablet(ctx, t, ts, tablet, nil)
 
 	resp, err := vtctld.GetTablet(ctx, &vtctldatapb.GetTabletRequest{
 		TabletAlias: &topodatapb.TabletAlias{
@@ -2198,14 +2194,14 @@ func TestGetSchema(t *testing.T) {
 	}
 	testutil.AddTablet(ctx, t, ts, &topodatapb.Tablet{
 		Alias: validAlias,
-	})
+	}, nil)
 	otherAlias := &topodatapb.TabletAlias{
 		Cell: "zone1",
 		Uid:  101,
 	}
 	testutil.AddTablet(ctx, t, ts, &topodatapb.Tablet{
 		Alias: otherAlias,
-	})
+	}, nil)
 
 	// we need to run this on each test case or they will pollute each other
 	setupSchema := func() {
@@ -2869,9 +2865,7 @@ func TestGetTablets(t *testing.T) {
 			ts := memorytopo.NewServer(tt.cells...)
 			vtctld := NewVtctldServer(ts)
 
-			for _, tablet := range tt.tablets {
-				testutil.AddTablet(ctx, t, ts, tablet)
-			}
+			testutil.AddTablets(ctx, t, ts, nil, tt.tablets...)
 
 			resp, err := vtctld.GetTablets(ctx, tt.req)
 			if tt.shouldErr {

go/vt/vtctl/grpcvtctldserver/testutil/util.go

@@ -20,13 +20,15 @@ package testutil
 
 import (
 	"context"
+	"fmt"
 	"testing"
 
 	"github.com/stretchr/testify/require"
 	"golang.org/x/net/nettest"
 	"google.golang.org/grpc"
 
 	"vitess.io/vitess/go/vt/topo"
+	"vitess.io/vitess/go/vt/topo/topoproto"
 	"vitess.io/vitess/go/vt/vtctl/vtctldclient"
 
 	topodatapb "vitess.io/vitess/go/vt/proto/topodata"
@@ -77,21 +79,48 @@ func AddKeyspaces(ctx context.Context, t *testing.T, ts *topo.Server, keyspaces
 	}
 }
 
+// AddTabletOptions is a container for different behaviors tests need from
+// AddTablet.
+type AddTabletOptions struct {
+	// AlsoSetShardMaster is an option to control additional setup to take when
+	// AddTablet receives a tablet of type MASTER. When set, AddTablet will also
+	// update the shard record to make that tablet the primary, and fail the
+	// test if the shard record has a serving primary already.
+	AlsoSetShardMaster bool
+	// SkipShardCreation, when set, makes AddTablet never attempt to create a
+	// shard record in the topo under any circumstances.
+	SkipShardCreation bool
+}
+
 // AddTablet adds a tablet to the topology, failing a test if that tablet record
 // could not be created. It shallow copies to prevent XXX_ fields from changing,
 // including nested proto message fields.
 //
 // AddTablet also optionally adds empty keyspace and shard records to the
 // topology, if they are set on the tablet record and they cannot be retrieved
 // from the topo server without error.
-func AddTablet(ctx context.Context, t *testing.T, ts *topo.Server, tablet *topodatapb.Tablet) {
+//
+// If AddTablet receives a tablet record with a keyspace and shard set, and that
+// tablet's type is MASTER, and opts.AlsoSetShardMaster is set, then AddTablet
+// will update the shard record to make that tablet the shard master and set the
+// shard to serving. If that shard record already has a serving primary, then
+// AddTablet will fail the test.
+func AddTablet(ctx context.Context, t *testing.T, ts *topo.Server, tablet *topodatapb.Tablet, opts *AddTabletOptions) {
 	in := *tablet
 	alias := *tablet.Alias
 	in.Alias = &alias
 
+	if opts == nil {
+		opts = &AddTabletOptions{}
+	}
+
 	err := ts.CreateTablet(ctx, &in)
 	require.NoError(t, err, "CreateTablet(%+v)", &in)
 
+	if opts.SkipShardCreation {
+		return
+	}
+
 	if tablet.Keyspace != "" {
 		if _, err := ts.GetKeyspace(ctx, tablet.Keyspace); err != nil {
 			err := ts.CreateKeyspace(ctx, tablet.Keyspace, &topodatapb.Keyspace{})
@@ -103,15 +132,30 @@ func AddTablet(ctx context.Context, t *testing.T, ts *topo.Server, tablet *topod
 				err := ts.CreateShard(ctx, tablet.Keyspace, tablet.Shard)
 				require.NoError(t, err, "CreateShard(%s, %s)", tablet.Keyspace, tablet.Shard)
 			}
+
+			if tablet.Type == topodatapb.TabletType_MASTER && opts.AlsoSetShardMaster {
+				_, err := ts.UpdateShardFields(ctx, tablet.Keyspace, tablet.Shard, func(si *topo.ShardInfo) error {
+					if si.IsMasterServing && si.MasterAlias != nil {
+						return fmt.Errorf("shard %v/%v already has a serving master (%v)", tablet.Keyspace, tablet.Shard, topoproto.TabletAliasString(si.MasterAlias))
+					}
+
+					si.MasterAlias = tablet.Alias
+					si.IsMasterServing = true
+					si.MasterTermStartTime = tablet.MasterTermStartTime
+
+					return nil
+				})
+				require.NoError(t, err, "UpdateShardFields(%s, %s) to set %s as serving primary failed", tablet.Keyspace, tablet.Shard, topoproto.TabletAliasString(tablet.Alias))
+			}
 		}
 	}
 }
 
 // AddTablets adds a list of tablets to the topology. See AddTablet for more
 // details.
-func AddTablets(ctx context.Context, t *testing.T, ts *topo.Server, tablets ...*topodatapb.Tablet) {
+func AddTablets(ctx context.Context, t *testing.T, ts *topo.Server, opts *AddTabletOptions, tablets ...*topodatapb.Tablet) {
 	for _, tablet := range tablets {
-		AddTablet(ctx, t, ts, tablet)
+		AddTablet(ctx, t, ts, tablet, opts)
 	}
 }
 
@@ -130,6 +174,15 @@ func AddShards(ctx context.Context, t *testing.T, ts *topo.Server, shards ...*vt
 
 		err := ts.CreateShard(ctx, shard.Keyspace, shard.Name)
 		require.NoError(t, err, "CreateShard(%s/%s)", shard.Keyspace, shard.Name)
+
+		if shard.Shard != nil {
+			_, err := ts.UpdateShardFields(ctx, shard.Keyspace, shard.Name, func(si *topo.ShardInfo) error {
+				si.Shard = shard.Shard
+
+				return nil
+			})
+			require.NoError(t, err, "UpdateShardFields(%s/%s, %v)", shard.Keyspace, shard.Name, shard.Shard)
+		}
 	}
 }