Add routing rules helper to test framework

Signed-off-by: Matt Lord <mattalord@gmail.com>

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

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)