server: call WaitFor5NodeReplication in TestCheckRestartSafe

This test should call the more specialized method to avoid timing out
when we're waiting for 5 nodes to replicate.

Epic: None
Resolves: #153523
Release note: None

Author: dhartunian

pkg/server/api_v2_test.go

@@ -444,7 +444,11 @@ func updateAllReplicaCounts(
 ) {
 	_, err := testCluster.ServerConn(0).Exec("ALTER RANGE default CONFIGURE ZONE USING num_replicas=$1", desiredVoters)
 	require.NoError(t, err)
-	require.NoError(t, testCluster.WaitForFullReplication())
+	if desiredVoters == 5 {
+		require.NoError(t, testCluster.WaitFor5NodeReplication())
+	} else {
+		require.NoError(t, testCluster.WaitForFullReplication())
+	}
 
 	ts0 := testCluster.Server(0)
 	stores := ts0.GetStores().(storeVisitor)

pkg/testutils/serverutils/test_cluster_shim.go

@@ -239,6 +239,11 @@ type TestClusterInterface interface {
 	// have no ranges with replication pending.
 	WaitForFullReplication() error
 
+	// WaitFor5NodeReplication ensures that zone configs are applied and
+	// up-replication is performed with new zone configs. This is the case for 5+
+	// node clusters.
+	WaitFor5NodeReplication() error
+
 	// StartedDefaultTestTenant returns whether this cluster started a
 	// default tenant for testing.
 	StartedDefaultTestTenant() bool