From dddb93585d7a34a602698729662ac1213e6f17d2 Mon Sep 17 00:00:00 2001
From: Ryan Leung <rleungx@gmail.com>
Date: Mon, 2 Dec 2024 15:18:38 +0800
Subject: [PATCH] fix tests

Signed-off-by: Ryan Leung <rleungx@gmail.com>
---
 server/cluster/cluster.go    | 40 +++++++++++++++++++++---------------
 tests/server/tso/tso_test.go |  1 +
 2 files changed, 25 insertions(+), 16 deletions(-)

diff --git a/server/cluster/cluster.go b/server/cluster/cluster.go
index 3d0e644c22c..c6542cb8511 100644
--- a/server/cluster/cluster.go
+++ b/server/cluster/cluster.go
@@ -419,32 +419,40 @@ func (c *RaftCluster) checkSchedulingService() {
 // checkTSOService checks the TSO service.
 func (c *RaftCluster) checkTSOService() {
 	if !c.opt.GetMicroServiceConfig().IsTSODynamicSwitchingEnabled() {
-		if err := c.startTSOJobsIfNeeded(); err != nil {
+		if err := c.switchToInternalTSO(); err != nil {
 			log.Error("failed to start TSO jobs", errs.ZapError(err))
 			return
 		}
-		if c.IsServiceIndependent(constant.TSOServiceName) {
-			c.UnsetServiceIndependent(constant.TSOServiceName)
-		}
 		return
 	}
 
 	servers, err := discovery.Discover(c.etcdClient, constant.TSOServiceName)
 	if err != nil || len(servers) == 0 {
-		if err := c.startTSOJobsIfNeeded(); err != nil {
-			log.Error("failed to start TSO jobs", errs.ZapError(err))
+		if err := c.switchToInternalTSO(); err != nil {
+			log.Error("failed to switch to internal TSO", errs.ZapError(err))
 			return
 		}
-		if c.IsServiceIndependent(constant.TSOServiceName) {
-			log.Info("TSO is provided by PD")
-			c.UnsetServiceIndependent(constant.TSOServiceName)
-		}
-	} else {
-		c.stopTSOJobsIfNeeded()
-		if !c.IsServiceIndependent(constant.TSOServiceName) {
-			log.Info("TSO is provided by TSO server")
-			c.SetServiceIndependent(constant.TSOServiceName)
-		}
+	} else if len(servers) > 0 {
+		c.switchToExternalTSO()
+	}
+}
+
+func (c *RaftCluster) switchToInternalTSO() error {
+	if err := c.startTSOJobsIfNeeded(); err != nil {
+		return err
+	}
+	if c.IsServiceIndependent(constant.TSOServiceName) {
+		c.UnsetServiceIndependent(constant.TSOServiceName)
+		log.Info("successfully switched to internal TSO")
+	}
+	return nil
+}
+
+func (c *RaftCluster) switchToExternalTSO() {
+	c.stopTSOJobsIfNeeded()
+	if !c.IsServiceIndependent(constant.TSOServiceName) {
+		c.SetServiceIndependent(constant.TSOServiceName)
+		log.Info("successfully switched to external TSO")
 	}
 }
 
diff --git a/tests/server/tso/tso_test.go b/tests/server/tso/tso_test.go
index cb6b87c83d3..7e4ec1276a0 100644
--- a/tests/server/tso/tso_test.go
+++ b/tests/server/tso/tso_test.go
@@ -132,6 +132,7 @@ func checkAndReturnTimestampResponse(re *require.Assertions, req *pdpb.TsoReques
 	re.GreaterOrEqual(uint32(timestamp.GetLogical())>>timestamp.GetSuffixBits(), req.GetCount())
 	return timestamp
 }
+
 func TestLogicalOverflow(t *testing.T) {
 	re := require.New(t)