Fix the tests

Signed-off-by: JmPotato <ghzpotato@gmail.com>
tikv · Jul 10, 2023 · d9c6b5e · d9c6b5e
1 parent 1d55fa5
commit d9c6b5e
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 13 deletions.
diff --git a/tests/pdctl/keyspace/keyspace_group_test.go b/tests/pdctl/keyspace/keyspace_group_test.go
@@ -203,10 +203,7 @@ func TestSetNodeAndPriorityKeyspaceGroup(t *testing.T) {
 	ttc, err := tests.NewTestTSOCluster(ctx, 2, pdAddr)
 	re.NoError(err)
 	defer ttc.Destroy()
-	tsoAddrs := make([]string, 0, 2)
-	for _, server := range ttc.GetServers() {
-		tsoAddrs = append(tsoAddrs, server.GetAddr())
-	}
+	tsoAddrs := ttc.GetAddrs()
 	cmd := pdctlCmd.GetRootCmd()
 
 	tc.WaitLeader()
@@ -514,12 +511,10 @@ func TestShowKeyspaceGroupPrimary(t *testing.T) {
 	re.NoError(err)
 	pdAddr := tc.GetConfig().GetClientURL()
 
-	s1, tsoServerCleanup1, err := tests.StartSingleTSOTestServer(ctx, re, pdAddr, tempurl.Alloc())
-	defer tsoServerCleanup1()
-	re.NoError(err)
-	s2, tsoServerCleanup2, err := tests.StartSingleTSOTestServer(ctx, re, pdAddr, tempurl.Alloc())
-	defer tsoServerCleanup2()
+	ttc, err := tests.NewTestTSOCluster(ctx, 2, pdAddr)
 	re.NoError(err)
+	defer ttc.Destroy()
+	tsoAddrs := ttc.GetAddrs()
 	cmd := pdctlCmd.GetRootCmd()
 
 	tc.WaitLeader()
@@ -540,7 +535,7 @@ func TestShowKeyspaceGroupPrimary(t *testing.T) {
 		return len(keyspaceGroup.Members) == 2
 	})
 	for _, member := range keyspaceGroup.Members {
-		re.Contains([]string{s1.GetAddr(), s2.GetAddr()}, member.Address)
+		re.Contains(tsoAddrs, member.Address)
 	}
 
 	// get primary for keyspace group 0.
@@ -550,7 +545,7 @@ func TestShowKeyspaceGroupPrimary(t *testing.T) {
 		re.NoError(err)
 		var resp handlers.GetKeyspaceGroupPrimaryResponse
 		json.Unmarshal(output, &resp)
-		return s1.GetAddr() == resp.Primary || s2.GetAddr() == resp.Primary
+		return tsoAddrs[0] == resp.Primary || tsoAddrs[1] == resp.Primary
 	})
 
 	// split keyspace group.
@@ -572,7 +567,7 @@ func TestShowKeyspaceGroupPrimary(t *testing.T) {
 		return len(keyspaceGroup.Members) == 2
 	})
 	for _, member := range keyspaceGroup.Members {
-		re.Contains([]string{s1.GetAddr(), s2.GetAddr()}, member.Address)
+		re.Contains(tsoAddrs, member.Address)
 	}
 
 	// get primary for keyspace group 1.
@@ -582,7 +577,7 @@ func TestShowKeyspaceGroupPrimary(t *testing.T) {
 		re.NoError(err)
 		var resp handlers.GetKeyspaceGroupPrimaryResponse
 		json.Unmarshal(output, &resp)
-		return s1.GetAddr() == resp.Primary || s2.GetAddr() == resp.Primary
+		return tsoAddrs[0] == resp.Primary || tsoAddrs[1] == resp.Primary
 	})
 
 	re.NoError(failpoint.Disable("github.com/tikv/pd/pkg/keyspace/acceleratedAllocNodes"))

diff --git a/tests/tso_cluster.go b/tests/tso_cluster.go
@@ -208,3 +208,12 @@ func (tc *TestTSOCluster) GetKeyspaceGroupMember() (members []endpoint.KeyspaceG
 	}
 	return
 }
+
+// GetAddrs returns all TSO server addresses.
+func (tc *TestTSOCluster) GetAddrs() []string {
+	addrs := make([]string, 0, len(tc.servers))
+	for _, server := range tc.servers {
+		addrs = append(addrs, server.GetAddr())
+	}
+	return addrs
+}