diff --git a/tests/integrations/mcs/tso/keyspace_group_manager_test.go b/tests/integrations/mcs/tso/keyspace_group_manager_test.go index 31286c8e190a..799fccd42e48 100644 --- a/tests/integrations/mcs/tso/keyspace_group_manager_test.go +++ b/tests/integrations/mcs/tso/keyspace_group_manager_test.go @@ -98,20 +98,22 @@ func (suite *tsoKeyspaceGroupManagerTestSuite) TestKeyspacesServedByDefaultKeysp // a keyspace group before, will be served by the default keyspace group. re := suite.Require() testutil.Eventually(re, func() bool { - for _, server := range suite.tsoCluster.GetServers() { - allServed := true - for _, keyspaceID := range []uint32{0, 1, 2} { + for _, keyspaceID := range []uint32{0, 1, 2} { + served := false + for _, server := range suite.tsoCluster.GetServers() { if server.IsKeyspaceServing(keyspaceID, mcsutils.DefaultKeyspaceGroupID) { tam, err := server.GetTSOAllocatorManager(mcsutils.DefaultKeyspaceGroupID) re.NoError(err) re.NotNil(tam) - } else { - allServed = false + served = true + break } } - return allServed + if !served { + return false + } } - return false + return true }, testutil.WithWaitFor(5*time.Second), testutil.WithTickInterval(50*time.Millisecond)) // Any keyspace that was assigned to a keyspace group before, except default keyspace, diff --git a/tests/integrations/tso/client_test.go b/tests/integrations/tso/client_test.go index 055490a159f4..bcf9456d8753 100644 --- a/tests/integrations/tso/client_test.go +++ b/tests/integrations/tso/client_test.go @@ -137,6 +137,24 @@ func (suite *tsoClientTestSuite) SetupSuite() { suite.keyspaceIDs = append(suite.keyspaceIDs, keyspaceGroup.keyspaceIDs...) } + // Make sure all keyspace groups are available. + testutil.Eventually(re, func() bool { + for _, keyspaceID := range suite.keyspaceIDs { + served := false + for _, server := range suite.tsoCluster.GetServers() { + if server.IsKeyspaceServing(keyspaceID, mcsutils.DefaultKeyspaceGroupID) { + served = true + break + } + } + if !served { + return false + } + } + return true + }, testutil.WithWaitFor(5*time.Second), testutil.WithTickInterval(50*time.Millisecond)) + + // Create clients and make sure they all have discovered the tso service. suite.clients = mcs.WaitForMultiKeyspacesTSOAvailable( suite.ctx, re, suite.keyspaceIDs, strings.Split(suite.backendEndpoints, ",")) re.Equal(len(suite.keyspaceIDs), len(suite.clients))