test: fix flakiness in move tables by making copier select the replic…

…a tablet

Signed-off-by: Manan Gupta <[email protected]>

go/test/endtoend/cluster/move_tables.go

@@ -18,6 +18,7 @@ package cluster
 
 import (
 	"fmt"
+	"strings"
 	"testing"
 	"time"
 )
@@ -31,17 +32,19 @@ type MoveTablesWorkflow struct {
 	targetKs        string
 	srcKs           string
 	tables          string
+	tabletTypes     []string
 }
 
 // NewMoveTables creates a new MoveTablesWorkflow.
-func NewMoveTables(t *testing.T, clusterInstance *LocalProcessCluster, workflowName, targetKs, srcKs, tables string) *MoveTablesWorkflow {
+func NewMoveTables(t *testing.T, clusterInstance *LocalProcessCluster, workflowName, targetKs, srcKs, tables string, tabletTypes []string) *MoveTablesWorkflow {
 	return &MoveTablesWorkflow{
 		t:               t,
 		clusterInstance: clusterInstance,
 		workflowName:    workflowName,
 		tables:          tables,
 		targetKs:        targetKs,
 		srcKs:           srcKs,
+		tabletTypes:     tabletTypes,
 	}
 }
 
@@ -52,6 +55,10 @@ func (mtw *MoveTablesWorkflow) Create() (string, error) {
 	} else {
 		args = append(args, "--all-tables")
 	}
+	if len(mtw.tabletTypes) != 0 {
+		args = append(args, "--tablet-types")
+		args = append(args, strings.Join(mtw.tabletTypes, ","))
+	}
 	return mtw.exec(args...)
 }
 

go/test/endtoend/transaction/twopc/stress/fuzzer_test.go

@@ -462,7 +462,7 @@ func moveTablesFuzzer(t *testing.T) {
 		require.NoError(t, err)
 	}
 	log.Errorf("MoveTables from - %v to %v", srcKeyspace, targetKeyspace)
-	mtw := cluster.NewMoveTables(t, clusterInstance, workflow, targetKeyspace, srcKeyspace, "twopc_fuzzer_update")
+	mtw := cluster.NewMoveTables(t, clusterInstance, workflow, targetKeyspace, srcKeyspace, "twopc_fuzzer_update", []string{"REPLICA"})
 	// Initiate MoveTables for twopc_fuzzer_update.
 	output, err := mtw.Create()
 	if err != nil {

go/test/endtoend/transaction/twopc/stress/stress_test.go

@@ -258,7 +258,7 @@ func mysqlRestartShard3(t *testing.T) error {
 // moveTablesCancel runs a move tables command that we cancel in the end.
 func moveTablesCancel(t *testing.T) error {
 	workflow := "TestDisruptions"
-	mtw := cluster.NewMoveTables(t, clusterInstance, workflow, unshardedKeyspaceName, keyspaceName, "twopc_t1")
+	mtw := cluster.NewMoveTables(t, clusterInstance, workflow, unshardedKeyspaceName, keyspaceName, "twopc_t1", []string{"REPLICA"})
 	// Initiate MoveTables for twopc_t1.
 	output, err := mtw.Create()
 	require.NoError(t, err, output)
@@ -277,7 +277,7 @@ func moveTablesCancel(t *testing.T) error {
 // moveTablesComplete runs a move tables command that we complete in the end.
 func moveTablesComplete(t *testing.T) error {
 	workflow := "TestDisruptions"
-	mtw := cluster.NewMoveTables(t, clusterInstance, workflow, unshardedKeyspaceName, keyspaceName, "twopc_t1")
+	mtw := cluster.NewMoveTables(t, clusterInstance, workflow, unshardedKeyspaceName, keyspaceName, "twopc_t1", []string{"REPLICA"})
 	// Initiate MoveTables for twopc_t1.
 	output, err := mtw.Create()
 	require.NoError(t, err, output)
@@ -297,7 +297,7 @@ func moveTablesReset(t *testing.T) {
 	err := clusterInstance.VtctldClientProcess.ApplyVSchema(keyspaceName, VSchema)
 	require.NoError(t, err)
 	workflow := "TestDisruptions"
-	mtw := cluster.NewMoveTables(t, clusterInstance, workflow, keyspaceName, unshardedKeyspaceName, "twopc_t1")
+	mtw := cluster.NewMoveTables(t, clusterInstance, workflow, keyspaceName, unshardedKeyspaceName, "twopc_t1", []string{"REPLICA"})
 	// Initiate MoveTables for twopc_t1.
 	output, err := mtw.Create()
 	require.NoError(t, err, output)

go/test/endtoend/transaction/twopc/utils/utils.go

@@ -27,6 +27,7 @@ import (
 	"github.com/stretchr/testify/require"
 
 	"vitess.io/vitess/go/mysql"
+	"vitess.io/vitess/go/vt/log"
 )
 
 const (
@@ -48,14 +49,14 @@ func ClearOutTable(t *testing.T, vtParams mysql.ConnParams, tableName string) {
 		}
 		conn, err := mysql.Connect(ctx, &vtParams)
 		if err != nil {
-			fmt.Printf("Error in connection - %v\n", err)
+			log.Errorf("Error in connection - %v\n", err)
 			time.Sleep(100 * time.Millisecond)
 			continue
 		}
 
 		res, err := conn.ExecuteFetch(fmt.Sprintf("SELECT count(*) FROM %v", tableName), 1, false)
 		if err != nil {
-			fmt.Printf("Error in selecting - %v\n", err)
+			log.Errorf("Error in selecting - %v\n", err)
 			conn.Close()
 			time.Sleep(100 * time.Millisecond)
 			continue
@@ -71,7 +72,7 @@ func ClearOutTable(t *testing.T, vtParams mysql.ConnParams, tableName string) {
 		_, err = conn.ExecuteFetch(fmt.Sprintf("DELETE FROM %v LIMIT 10000", tableName), 10000, false)
 		conn.Close()
 		if err != nil {
-			fmt.Printf("Error in cleanup deletion - %v\n", err)
+			log.Errorf("Error in cleanup deletion - %v\n", err)
 			time.Sleep(100 * time.Millisecond)
 			continue
 		}