VReplication workflow package: unit tests for StreamMigrator, Mount e…

…t al (vitessio#16498)

Signed-off-by: Rohit Nayak <rohit@planetscale.com>

go/vt/vtctl/workflow/framework_test.go

@@ -256,6 +256,9 @@ type testTMClient struct {
 	readVReplicationWorkflowRequests   map[uint32]*tabletmanagerdatapb.ReadVReplicationWorkflowRequest
 	primaryPositions                   map[uint32]string
 
+	// Stack of ReadVReplicationWorkflowsResponse to return, in order, for each shard
+	readVReplicationWorkflowsResponses map[string][]*tabletmanagerdatapb.ReadVReplicationWorkflowsResponse
+
 	env     *testEnv    // For access to the env config from tmc methods.
 	reverse atomic.Bool // Are we reversing traffic?
 	frozen  atomic.Bool // Are the workflows frozen?
@@ -267,6 +270,7 @@ func newTestTMClient(env *testEnv) *testTMClient {
 		vrQueries:                          make(map[int][]*queryResult),
 		createVReplicationWorkflowRequests: make(map[uint32]*tabletmanagerdatapb.CreateVReplicationWorkflowRequest),
 		readVReplicationWorkflowRequests:   make(map[uint32]*tabletmanagerdatapb.ReadVReplicationWorkflowRequest),
+		readVReplicationWorkflowsResponses: make(map[string][]*tabletmanagerdatapb.ReadVReplicationWorkflowsResponse),
 		primaryPositions:                   make(map[uint32]string),
 		env:                                env,
 	}
@@ -285,6 +289,10 @@ func (tmc *testTMClient) CreateVReplicationWorkflow(ctx context.Context, tablet
 	return &tabletmanagerdatapb.CreateVReplicationWorkflowResponse{Result: sqltypes.ResultToProto3(res)}, nil
 }
 
+func (tmc *testTMClient) GetWorkflowKey(keyspace, shard string) string {
+	return fmt.Sprintf("%s/%s", keyspace, shard)
+}
+
 func (tmc *testTMClient) ReadVReplicationWorkflow(ctx context.Context, tablet *topodatapb.Tablet, req *tabletmanagerdatapb.ReadVReplicationWorkflowRequest) (*tabletmanagerdatapb.ReadVReplicationWorkflowResponse, error) {
 	tmc.mu.Lock()
 	defer tmc.mu.Unlock()
@@ -463,6 +471,10 @@ func (tmc *testTMClient) ReadVReplicationWorkflows(ctx context.Context, tablet *
 	tmc.mu.Lock()
 	defer tmc.mu.Unlock()
 
+	workflowKey := tmc.GetWorkflowKey(tablet.Keyspace, tablet.Shard)
+	if resp := tmc.getVReplicationWorkflowsResponse(workflowKey); resp != nil {
+		return resp, nil
+	}
 	workflowType := binlogdatapb.VReplicationWorkflowType_MoveTables
 	if len(req.IncludeWorkflows) > 0 {
 		for _, wf := range req.IncludeWorkflows {
@@ -494,7 +506,7 @@ func (tmc *testTMClient) ReadVReplicationWorkflows(ctx context.Context, tablet *
 									},
 								},
 							},
-							Pos:           "MySQL56/" + position,
+							Pos:           position,
 							TimeUpdated:   protoutil.TimeToProto(time.Now()),
 							TimeHeartbeat: protoutil.TimeToProto(time.Now()),
 						},
@@ -541,6 +553,25 @@ func (tmc *testTMClient) VReplicationWaitForPos(ctx context.Context, tablet *top
 	return nil
 }
 
+func (tmc *testTMClient) AddVReplicationWorkflowsResponse(key string, resp *tabletmanagerdatapb.ReadVReplicationWorkflowsResponse) {
+	tmc.mu.Lock()
+	defer tmc.mu.Unlock()
+	tmc.readVReplicationWorkflowsResponses[key] = append(tmc.readVReplicationWorkflowsResponses[key], resp)
+}
+
+func (tmc *testTMClient) getVReplicationWorkflowsResponse(key string) *tabletmanagerdatapb.ReadVReplicationWorkflowsResponse {
+	if len(tmc.readVReplicationWorkflowsResponses) == 0 {
+		return nil
+	}
+	responses, ok := tmc.readVReplicationWorkflowsResponses[key]
+	if !ok || len(responses) == 0 {
+		return nil
+	}
+	resp := tmc.readVReplicationWorkflowsResponses[key][0]
+	tmc.readVReplicationWorkflowsResponses[key] = tmc.readVReplicationWorkflowsResponses[key][1:]
+	return resp
+}
+
 //
 // Utility / helper functions.
 //

go/vt/vtctl/workflow/materializer_env_test.go

@@ -61,7 +61,7 @@ type testMaterializerEnv struct {
 	venv     *vtenv.Environment
 }
 
-//----------------------------------------------
+// ----------------------------------------------
 // testMaterializerEnv
 
 func newTestMaterializerEnv(t *testing.T, ctx context.Context, ms *vtctldatapb.MaterializeSettings, sourceShards, targetShards []string) *testMaterializerEnv {
@@ -426,7 +426,7 @@ func (tmc *testMaterializerTMClient) ReadVReplicationWorkflows(ctx context.Conte
 						},
 					},
 				},
-				Pos:           "MySQL56/" + position,
+				Pos:           position,
 				TimeUpdated:   protoutil.TimeToProto(time.Now()),
 				TimeHeartbeat: protoutil.TimeToProto(time.Now()),
 			}

go/vt/vtctl/workflow/materializer_test.go

@@ -44,7 +44,7 @@ import (
 )
 
 const (
-	position             = "9d10e6ec-07a0-11ee-ae73-8e53f4cf3083:1-97"
+	position             = "MySQL56/9d10e6ec-07a0-11ee-ae73-8e53f4cf3083:1-97"
 	mzSelectFrozenQuery  = "select 1 from _vt.vreplication where db_name='vt_targetks' and message='FROZEN' and workflow_sub_type != 1"
 	mzCheckJournal       = "/select val from _vt.resharding_journal where id="
 	mzGetCopyState       = "select distinct table_name from _vt.copy_state cs, _vt.vreplication vr where vr.id = cs.vrepl_id and vr.id = 1"
@@ -56,6 +56,14 @@ var (
 	defaultOnDDL = binlogdatapb.OnDDLAction_IGNORE.String()
 )
 
+func gtid(position string) string {
+	arr := strings.Split(position, "/")
+	if len(arr) != 2 {
+		return ""
+	}
+	return arr[1]
+}
+
 func TestStripForeignKeys(t *testing.T) {
 	tcs := []struct {
 		desc string
@@ -577,7 +585,7 @@ func TestMoveTablesDDLFlag(t *testing.T) {
 			sourceShard, err := env.topoServ.GetShardNames(ctx, ms.SourceKeyspace)
 			require.NoError(t, err)
 			want := fmt.Sprintf("shard_streams:{key:\"%s/%s\" value:{streams:{id:1 tablet:{cell:\"%s\" uid:200} source_shard:\"%s/%s\" position:\"%s\" status:\"Running\" info:\"VStream Lag: 0s\"}}} traffic_state:\"Reads Not Switched. Writes Not Switched\"",
-				ms.TargetKeyspace, targetShard[0], env.cell, ms.SourceKeyspace, sourceShard[0], position)
+				ms.TargetKeyspace, targetShard[0], env.cell, ms.SourceKeyspace, sourceShard[0], gtid(position))
 
 			res, err := env.ws.MoveTablesCreate(ctx, &vtctldatapb.MoveTablesCreateRequest{
 				Workflow:       ms.Workflow,
@@ -636,7 +644,7 @@ func TestMoveTablesNoRoutingRules(t *testing.T) {
 							Uid:  200,
 						},
 						SourceShard: fmt.Sprintf("%s/%s", ms.SourceKeyspace, sourceShard[0]),
-						Position:    position,
+						Position:    gtid(position),
 						Status:      binlogdatapb.VReplicationWorkflowState_Running.String(),
 						Info:        "VStream Lag: 0s",
 					},

go/vt/vtctl/workflow/mount_test.go

@@ -0,0 +1,77 @@
+/*
+Copyright 2024 The Vitess Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+	http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package workflow
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	vtctldatapb "vitess.io/vitess/go/vt/proto/vtctldata"
+	"vitess.io/vitess/go/vt/topo/memorytopo"
+	"vitess.io/vitess/go/vt/vtenv"
+)
+
+// TestMount tests various Mount-related methods.
+func TestMount(t *testing.T) {
+	const (
+		extCluster = "extcluster"
+		topoType   = "etcd2"
+		topoServer = "localhost:2379"
+		topoRoot   = "/vitess/global"
+	)
+	ctx := context.Background()
+	ts := memorytopo.NewServer(ctx, "cell")
+	tmc := &fakeTMC{}
+	s := NewServer(vtenv.NewTestEnv(), ts, tmc)
+
+	resp, err := s.MountRegister(ctx, &vtctldatapb.MountRegisterRequest{
+		Name:       extCluster,
+		TopoType:   topoType,
+		TopoServer: topoServer,
+		TopoRoot:   topoRoot,
+	})
+	require.NoError(t, err)
+	require.NotNil(t, resp)
+
+	respList, err := s.MountList(ctx, &vtctldatapb.MountListRequest{})
+	require.NoError(t, err)
+	require.NotNil(t, respList)
+	require.Equal(t, []string{extCluster}, respList.Names)
+
+	respShow, err := s.MountShow(ctx, &vtctldatapb.MountShowRequest{
+		Name: extCluster,
+	})
+	require.NoError(t, err)
+	require.NotNil(t, respShow)
+	require.Equal(t, extCluster, respShow.Name)
+	require.Equal(t, topoType, respShow.TopoType)
+	require.Equal(t, topoServer, respShow.TopoServer)
+	require.Equal(t, topoRoot, respShow.TopoRoot)
+
+	respUnregister, err := s.MountUnregister(ctx, &vtctldatapb.MountUnregisterRequest{
+		Name: extCluster,
+	})
+	require.NoError(t, err)
+	require.NotNil(t, respUnregister)
+
+	respList, err = s.MountList(ctx, &vtctldatapb.MountListRequest{})
+	require.NoError(t, err)
+	require.NotNil(t, respList)
+	require.Nil(t, respList.Names)
+}

go/vt/vtctl/workflow/resharder_test.go

@@ -84,7 +84,7 @@ func TestReshardCreate(t *testing.T) {
 							{
 								Id:          1,
 								Tablet:      &topodatapb.TabletAlias{Cell: defaultCellName, Uid: startingTargetTabletUID},
-								SourceShard: "targetks/0", Position: position, Status: "Running", Info: "VStream Lag: 0s",
+								SourceShard: "targetks/0", Position: gtid(position), Status: "Running", Info: "VStream Lag: 0s",
 							},
 						},
 					},
@@ -93,7 +93,7 @@ func TestReshardCreate(t *testing.T) {
 							{
 								Id:          1,
 								Tablet:      &topodatapb.TabletAlias{Cell: defaultCellName, Uid: startingTargetTabletUID + tabletUIDStep},
-								SourceShard: "targetks/0", Position: position, Status: "Running", Info: "VStream Lag: 0s",
+								SourceShard: "targetks/0", Position: gtid(position), Status: "Running", Info: "VStream Lag: 0s",
 							},
 						},
 					},