testutils/testcluster: automatically wait-for-full-replication

Closes cockroachdb#9624

pkg/base/test_server_args.go

@@ -104,8 +104,6 @@ const (
 	// ReplicationAuto means that ranges are replicated according to the
 	// production default zone config. Replication is performed as in
 	// production, by the replication queue.
-	// TestCluster.WaitForFullReplication() can be used to wait for
-	// replication to be stable at any point in a test.
 	ReplicationAuto TestClusterReplicationMode = iota
 	// ReplicationManual means that the split and replication queues of all
 	// servers are stopped, and the test must manually control splitting and

pkg/server/admin_cluster_test.go

@@ -42,9 +42,6 @@ func TestAdminAPITableStats(t *testing.T) {
 		},
 	})
 	defer tc.Stopper().Stop()
-	if err := tc.WaitForFullReplication(); err != nil {
-		t.Fatal(err)
-	}
 	server0 := tc.Server(0)
 
 	// Create clients (SQL, HTTP) connected to server 0.

pkg/sql/bench_test.go

@@ -59,9 +59,6 @@ func benchmarkMultinodeCockroach(b *testing.B, f func(b *testing.B, db *gosql.DB
 	if _, err := tc.Conns[0].Exec(`CREATE DATABASE bench`); err != nil {
 		b.Fatal(err)
 	}
-	if err := tc.WaitForFullReplication(); err != nil {
-		b.Fatal(err)
-	}
 	defer tc.Stopper().Stop()
 
 	f(b, tc.Conns[0])

pkg/sql/create_test.go

@@ -246,9 +246,6 @@ func TestParallelCreateTables(t *testing.T) {
 
 	tc := testcluster.StartTestCluster(t, numberOfNodes, base.TestClusterArgs{})
 	defer tc.Stopper().Stop()
-	if err := tc.WaitForFullReplication(); err != nil {
-		t.Fatal(err)
-	}
 
 	if _, err := tc.ServerConn(0).Exec(`CREATE DATABASE "test"`); err != nil {
 		t.Fatal(err)
@@ -302,9 +299,6 @@ func TestParallelCreateConflictingTables(t *testing.T) {
 
 	tc := testcluster.StartTestCluster(t, numberOfNodes, base.TestClusterArgs{})
 	defer tc.Stopper().Stop()
-	if err := tc.WaitForFullReplication(); err != nil {
-		t.Fatal(err)
-	}
 
 	if _, err := tc.ServerConn(0).Exec(`CREATE DATABASE "test"`); err != nil {
 		t.Fatal(err)

pkg/sql/parallel_test.go

@@ -225,14 +225,6 @@ func (t *parallelTest) setup(spec *parTestSpec) {
 		sqlutils.MakeSQLRunner(t, t.clients[i][0]).Exec("SET DATABASE = test")
 	}
 
-	if spec.ClusterSize >= 3 {
-		if testing.Verbose() || log.V(1) {
-			log.Infof(t.ctx, "Waiting for full replication")
-		}
-		if err := t.cluster.WaitForFullReplication(); err != nil {
-			t.Fatal(err)
-		}
-	}
 	if testing.Verbose() || log.V(1) {
 		log.Infof(t.ctx, "Test setup done")
 	}

pkg/sql/table_split_test.go

@@ -39,9 +39,6 @@ func TestSplitAtTableBoundary(t *testing.T) {
 	}
 	tc := testcluster.StartTestCluster(t, 3, testClusterArgs)
 	defer tc.Stopper().Stop()
-	if err := tc.WaitForFullReplication(); err != nil {
-		t.Fatal(err)
-	}
 
 	runner := sqlutils.MakeSQLRunner(t, tc.Conns[0])
 	runner.Exec(`CREATE DATABASE test`)

pkg/testutils/serverutils/test_cluster_shim.go

@@ -42,8 +42,6 @@ type TestClusterInterface interface {
 	// ServerConn returns a gosql.DB connection to a specific node.
 	ServerConn(idx int) *gosql.DB
 
-	WaitForFullReplication() error
-
 	// StopServer stops a single server.
 	StopServer(idx int)
 

pkg/testutils/testcluster/testcluster.go

@@ -178,6 +178,12 @@ func StartTestCluster(t testing.TB, nodes int, args base.TestClusterArgs) *TestC
 	tc.stopper.AddCloser(stop.CloserFn(tc.stopServers))
 
 	tc.waitForStores(t)
+
+	if args.ReplicationMode == base.ReplicationAuto && nodes >= 3 {
+		if err := tc.waitForFullReplication(); err != nil {
+			t.Fatal(err)
+		}
+	}
 	return tc
 }
 
@@ -490,9 +496,9 @@ func (tc *TestCluster) findMemberStore(storeID roachpb.StoreID) (*storage.Store,
 	return nil, errors.Errorf("store not found")
 }
 
-// WaitForFullReplication waits until all stores in the cluster
+// waitForFullReplication waits until all stores in the cluster
 // have no ranges with replication pending.
-func (tc *TestCluster) WaitForFullReplication() error {
+func (tc *TestCluster) waitForFullReplication() error {
 	opts := retry.Options{
 		InitialBackoff: time.Millisecond * 10,
 		MaxBackoff:     time.Millisecond * 100,

pkg/testutils/testcluster/testcluster_test.go

@@ -169,14 +169,12 @@ func TestBasicManualReplication(t *testing.T) {
 	}
 }
 
-func TestWaitForFullReplication(t *testing.T) {
+func TestBasicAutoReplication(t *testing.T) {
 	defer leaktest.AfterTest(t)()
 
 	tc := StartTestCluster(t, 3, base.TestClusterArgs{ReplicationMode: base.ReplicationAuto})
 	defer tc.Stopper().Stop()
-	if err := tc.WaitForFullReplication(); err != nil {
-		t.Error(err)
-	}
+	// NB: StartTestCluster will wait for full replication.
 }
 
 func TestStopServer(t *testing.T) {
@@ -192,9 +190,6 @@ func TestStopServer(t *testing.T) {
 		ReplicationMode: base.ReplicationAuto,
 	})
 	defer tc.Stopper().Stop()
-	if err := tc.WaitForFullReplication(); err != nil {
-		t.Fatal(err)
-	}
 
 	// Connect to server 1, ensure it is answering requests over HTTP and GRPC.
 	server1 := tc.Server(1)