diff --git a/go/vt/vtctl/workflow/framework_test.go b/go/vt/vtctl/workflow/framework_test.go index f05e1e201fa..73b34015338 100644 --- a/go/vt/vtctl/workflow/framework_test.go +++ b/go/vt/vtctl/workflow/framework_test.go @@ -36,6 +36,7 @@ import ( "vitess.io/vitess/go/vt/mysqlctl/tmutils" "vitess.io/vitess/go/vt/topo" "vitess.io/vitess/go/vt/topo/memorytopo" + "vitess.io/vitess/go/vt/topotools" "vitess.io/vitess/go/vt/vtenv" "vitess.io/vitess/go/vt/vterrors" "vitess.io/vitess/go/vt/vttablet/tmclient" @@ -150,6 +151,32 @@ func (env *testEnv) addTablet(t *testing.T, ctx context.Context, id int, keyspac return tablet } +// addTableRoutingRules adds routing rules from the test env's source keyspace to +// its target keyspace for the given tablet types and tables. +func (env *testEnv) addTableRoutingRules(t *testing.T, ctx context.Context, tabletTypes []topodatapb.TabletType, tables []string) { + ks := env.targetKeyspace.KeyspaceName + rules := make(map[string][]string, len(tables)*(len(tabletTypes)*3)) + for _, tabletType := range tabletTypes { + for _, tableName := range tables { + toTarget := []string{ks + "." + tableName} + tt := strings.ToLower(tabletType.String()) + if tabletType == topodatapb.TabletType_PRIMARY { + rules[tableName] = toTarget + rules[ks+"."+tableName] = toTarget + rules[env.sourceKeyspace.KeyspaceName+"."+tableName] = toTarget + } else { + rules[tableName+"@"+tt] = toTarget + rules[ks+"."+tableName+"@"+tt] = toTarget + rules[env.sourceKeyspace.KeyspaceName+"."+tableName+"@"+tt] = toTarget + } + } + } + err := topotools.SaveRoutingRules(ctx, env.ts, rules) + require.NoError(t, err) + err = env.ts.RebuildSrvVSchema(ctx, nil) + require.NoError(t, err) +} + func (env *testEnv) deleteTablet(tablet *topodatapb.Tablet) { _ = env.ts.DeleteTablet(context.Background(), tablet.Alias) delete(env.tablets[tablet.Keyspace], int(tablet.Alias.Uid)) diff --git a/go/vt/vtctl/workflow/server_test.go b/go/vt/vtctl/workflow/server_test.go index cee8c6ac247..fb432403155 100644 --- a/go/vt/vtctl/workflow/server_test.go +++ b/go/vt/vtctl/workflow/server_test.go @@ -34,7 +34,6 @@ import ( "vitess.io/vitess/go/vt/topo" "vitess.io/vitess/go/vt/topo/memorytopo" "vitess.io/vitess/go/vt/topo/topoproto" - "vitess.io/vitess/go/vt/topotools" "vitess.io/vitess/go/vt/vtenv" "vitess.io/vitess/go/vt/vttablet/tmclient" @@ -555,25 +554,7 @@ func TestMoveTablesTrafficSwitching(t *testing.T) { } else { env.tmc.reverse.Store(true) // Setup the routing rules as they would be after having previously done SwitchTraffic. - ks := env.targetKeyspace.KeyspaceName - toTarget := []string{ks + "." + tableName} - rules := make(map[string][]string) - for _, tabletType := range tabletTypes { - tt := strings.ToLower(tabletType.String()) - if tabletType == topodatapb.TabletType_PRIMARY { - rules[tableName] = toTarget - rules[ks+"."+tableName] = toTarget - rules[env.sourceKeyspace.KeyspaceName+"."+tableName] = toTarget - } else { - rules[tableName+"@"+tt] = toTarget - rules[ks+"."+tableName+"@"+tt] = toTarget - rules[env.sourceKeyspace.KeyspaceName+"."+tableName+"@"+tt] = toTarget - } - } - err := topotools.SaveRoutingRules(ctx, env.ts, rules) - require.NoError(t, err) - err = env.ts.RebuildSrvVSchema(ctx, nil) - require.NoError(t, err) + env.addTableRoutingRules(t, ctx, tabletTypes, []string{tableName}) env.tmc.expectVRQueryResultOnKeyspaceTablets(tc.sourceKeyspace.KeyspaceName, copyTableQR) for i := 0; i < len(tc.targetKeyspace.ShardNames); i++ { // Per stream env.tmc.expectVRQueryResultOnKeyspaceTablets(tc.sourceKeyspace.KeyspaceName, cutoverQR) @@ -783,27 +764,7 @@ func TestMoveTablesTrafficSwitchingDryRun(t *testing.T) { } else { env.tmc.reverse.Store(true) // Setup the routing rules as they would be after having previously done SwitchTraffic. - ks := env.targetKeyspace.KeyspaceName - toTarget := []string{ks + "." + table1Name} - rules := make(map[string][]string) - for _, tabletType := range tabletTypes { - for _, tableName := range tables { - tt := strings.ToLower(tabletType.String()) - if tabletType == topodatapb.TabletType_PRIMARY { - rules[tableName] = toTarget - rules[ks+"."+tableName] = toTarget - rules[env.sourceKeyspace.KeyspaceName+"."+tableName] = toTarget - } else { - rules[tableName+"@"+tt] = toTarget - rules[ks+"."+tableName+"@"+tt] = toTarget - rules[env.sourceKeyspace.KeyspaceName+"."+tableName+"@"+tt] = toTarget - } - } - } - err := topotools.SaveRoutingRules(ctx, env.ts, rules) - require.NoError(t, err) - err = env.ts.RebuildSrvVSchema(ctx, nil) - require.NoError(t, err) + env.addTableRoutingRules(t, ctx, tabletTypes, tables) env.tmc.expectVRQueryResultOnKeyspaceTablets(tc.sourceKeyspace.KeyspaceName, copyTableQR) for i := 0; i < len(tc.targetKeyspace.ShardNames); i++ { // Per stream env.tmc.expectVRQueryResultOnKeyspaceTablets(tc.targetKeyspace.KeyspaceName, journalQR)