*: fix TestGetTSOImmediately (tikv#6811)

close tikv#6795

Signed-off-by: Ryan Leung <[email protected]>

tests/integrations/mcs/tso/keyspace_group_manager_test.go

@@ -17,6 +17,7 @@ package tso
 import (
 	"context"
 	"math/rand"
+	"net/http"
 	"strings"
 	"sync"
 	"testing"
@@ -357,13 +358,19 @@ func waitFinishSplit(
 	splitSourceKeyspaces, splitTargetKeyspaces []uint32,
 ) {
 	testutil.Eventually(re, func() bool {
-		kg := handlersutil.MustLoadKeyspaceGroupByID(re, server, splitTargetID)
+		kg, code := handlersutil.TryLoadKeyspaceGroupByID(re, server, splitTargetID)
+		if code != http.StatusOK {
+			return false
+		}
 		re.Equal(splitTargetID, kg.ID)
 		re.Equal(splitTargetKeyspaces, kg.Keyspaces)
 		return !kg.IsSplitTarget()
 	})
 	testutil.Eventually(re, func() bool {
-		kg := handlersutil.MustLoadKeyspaceGroupByID(re, server, splitSourceID)
+		kg, code := handlersutil.TryLoadKeyspaceGroupByID(re, server, splitSourceID)
+		if code != http.StatusOK {
+			return false
+		}
 		re.Equal(splitSourceID, kg.ID)
 		re.Equal(splitSourceKeyspaces, kg.Keyspaces)
 		return !kg.IsSplitSource()

tests/server/apiv2/handlers/testutil.go

@@ -168,17 +168,27 @@ func tryCreateKeyspaceGroup(re *require.Assertions, server *tests.TestServer, re
 
 // MustLoadKeyspaceGroupByID loads the keyspace group by ID with HTTP API.
 func MustLoadKeyspaceGroupByID(re *require.Assertions, server *tests.TestServer, id uint32) *endpoint.KeyspaceGroup {
+	kg, code := TryLoadKeyspaceGroupByID(re, server, id)
+	re.Equal(http.StatusOK, code)
+	return kg
+}
+
+// TryLoadKeyspaceGroupByID loads the keyspace group by ID with HTTP API.
+func TryLoadKeyspaceGroupByID(re *require.Assertions, server *tests.TestServer, id uint32) (*endpoint.KeyspaceGroup, int) {
 	httpReq, err := http.NewRequest(http.MethodGet, server.GetAddr()+keyspaceGroupsPrefix+fmt.Sprintf("/%d", id), nil)
 	re.NoError(err)
 	resp, err := dialClient.Do(httpReq)
 	re.NoError(err)
 	defer resp.Body.Close()
 	data, err := io.ReadAll(resp.Body)
 	re.NoError(err)
-	re.Equal(http.StatusOK, resp.StatusCode, string(data))
+	if resp.StatusCode != http.StatusOK {
+		return nil, resp.StatusCode
+	}
+
 	var kg endpoint.KeyspaceGroup
 	re.NoError(json.Unmarshal(data, &kg))
-	return &kg
+	return &kg, resp.StatusCode
 }
 
 // MustCreateKeyspaceGroup creates a keyspace group with HTTP API.