Make persistence tests consistent in handling ShardID (#3772)

Author: mxbvkn

common/persistence/tests/execution_mutable_state.go

@@ -86,8 +86,7 @@ func NewExecutionMutableStateSuite(
 			logger,
 			dynamicconfig.GetIntPropertyFn(4*1024*1024),
 		),
-		Logger:  logger,
-		ShardID: 1,
+		Logger: logger,
 	}
 }
 
@@ -103,7 +102,7 @@ func (s *ExecutionMutableStateSuite) SetupTest() {
 	s.Assertions = require.New(s.T())
 	s.Ctx, s.Cancel = context.WithTimeout(context.Background(), time.Second*30)
 
-	s.ShardID = 1 + s.ShardID
+	s.ShardID++
 	resp, err := s.ShardManager.GetOrCreateShard(s.Ctx, &p.GetOrCreateShardRequest{
 		ShardID: s.ShardID,
 		InitialShardInfo: &persistencespb.ShardInfo{
@@ -113,7 +112,7 @@ func (s *ExecutionMutableStateSuite) SetupTest() {
 	})
 	s.NoError(err)
 	previousRangeID := resp.ShardInfo.RangeId
-	resp.ShardInfo.RangeId += 1
+	resp.ShardInfo.RangeId++
 	err = s.ShardManager.UpdateShard(s.Ctx, &p.UpdateShardRequest{
 		ShardInfo:       resp.ShardInfo,
 		PreviousRangeID: previousRangeID,

common/persistence/tests/execution_mutable_state_task.go

@@ -96,16 +96,15 @@ func NewExecutionMutableStateTaskSuite(
 			logger,
 			dynamicconfig.GetIntPropertyFn(4*1024*1024),
 		),
-		Logger:  logger,
-		ShardID: 1,
+		Logger: logger,
 	}
 }
 
 func (s *ExecutionMutableStateTaskSuite) SetupTest() {
 	s.Assertions = require.New(s.T())
 	s.Ctx, s.Cancel = context.WithTimeout(context.Background(), time.Second*30)
 
-	s.ShardID = 1 + s.ShardID
+	s.ShardID++
 	resp, err := s.ShardManager.GetOrCreateShard(s.Ctx, &p.GetOrCreateShardRequest{
 		ShardID: s.ShardID,
 		InitialShardInfo: &persistencespb.ShardInfo{
@@ -115,7 +114,7 @@ func (s *ExecutionMutableStateTaskSuite) SetupTest() {
 	})
 	s.NoError(err)
 	previousRangeID := resp.ShardInfo.RangeId
-	resp.ShardInfo.RangeId += 1
+	resp.ShardInfo.RangeId++
 	err = s.ShardManager.UpdateShard(s.Ctx, &p.UpdateShardRequest{
 		ShardInfo:       resp.ShardInfo,
 		PreviousRangeID: previousRangeID,

common/persistence/tests/shard.go

@@ -43,6 +43,8 @@ type (
 		suite.Suite
 		*require.Assertions
 
+		ShardID int32
+
 		ShardManager p.ShardManager
 		Logger       log.Logger
 
@@ -78,19 +80,20 @@ func (s *ShardSuite) TearDownSuite() {
 func (s *ShardSuite) SetupTest() {
 	s.Assertions = require.New(s.T())
 	s.Ctx, s.Cancel = context.WithTimeout(context.Background(), time.Second*30)
+
+	s.ShardID++
 }
 
 func (s *ShardSuite) TearDownTest() {
 	s.Cancel()
 }
 
 func (s *ShardSuite) TestGetOrCreateShard_Create() {
-	shardID := int32(1)
 	rangeID := rand.Int63()
-	shardInfo := RandomShardInfo(shardID, rangeID)
+	shardInfo := RandomShardInfo(s.ShardID, rangeID)
 
 	resp, err := s.ShardManager.GetOrCreateShard(s.Ctx, &p.GetOrCreateShardRequest{
-		ShardID:          shardID,
+		ShardID:          s.ShardID,
 		InitialShardInfo: shardInfo,
 	})
 	s.NoError(err)
@@ -99,74 +102,71 @@ func (s *ShardSuite) TestGetOrCreateShard_Create() {
 }
 
 func (s *ShardSuite) TestGetOrCreateShard_Get() {
-	shardID := int32(2)
 	rangeID := rand.Int63()
-	shardInfo := RandomShardInfo(shardID, rangeID)
+	shardInfo := RandomShardInfo(s.ShardID, rangeID)
 
 	resp, err := s.ShardManager.GetOrCreateShard(s.Ctx, &p.GetOrCreateShardRequest{
-		ShardID:          shardID,
+		ShardID:          s.ShardID,
 		InitialShardInfo: shardInfo,
 	})
 	s.NoError(err)
 	s.Equal(shardInfo, resp.ShardInfo)
 
 	resp, err = s.ShardManager.GetOrCreateShard(s.Ctx, &p.GetOrCreateShardRequest{
-		ShardID:          shardID,
-		InitialShardInfo: RandomShardInfo(shardID, rand.Int63()),
+		ShardID:          s.ShardID,
+		InitialShardInfo: RandomShardInfo(s.ShardID, rand.Int63()),
 	})
 	s.NoError(err)
 	s.Equal(shardInfo, resp.ShardInfo)
 }
 
 func (s *ShardSuite) TestUpdateShard_OwnershipLost() {
-	shardID := int32(3)
 	rangeID := rand.Int63()
-	shardInfo := RandomShardInfo(shardID, rangeID)
+	shardInfo := RandomShardInfo(s.ShardID, rangeID)
 
 	resp, err := s.ShardManager.GetOrCreateShard(s.Ctx, &p.GetOrCreateShardRequest{
-		ShardID:          shardID,
+		ShardID:          s.ShardID,
 		InitialShardInfo: shardInfo,
 	})
 	s.NoError(err)
 	s.Equal(shardInfo, resp.ShardInfo)
 
 	updateRangeID := rand.Int63()
-	updateShardInfo := RandomShardInfo(shardID, rand.Int63())
+	updateShardInfo := RandomShardInfo(s.ShardID, rand.Int63())
 	err = s.ShardManager.UpdateShard(s.Ctx, &p.UpdateShardRequest{
 		ShardInfo:       updateShardInfo,
 		PreviousRangeID: updateRangeID,
 	})
 	s.IsType(&p.ShardOwnershipLostError{}, err)
 
 	resp, err = s.ShardManager.GetOrCreateShard(s.Ctx, &p.GetOrCreateShardRequest{
-		ShardID:          shardID,
+		ShardID:          s.ShardID,
 		InitialShardInfo: shardInfo,
 	})
 	s.NoError(err)
 	s.Equal(shardInfo, resp.ShardInfo)
 }
 
 func (s *ShardSuite) TestUpdateShard_Success() {
-	shardID := int32(4)
 	rangeID := rand.Int63()
-	shardInfo := RandomShardInfo(shardID, rangeID)
+	shardInfo := RandomShardInfo(s.ShardID, rangeID)
 
 	resp, err := s.ShardManager.GetOrCreateShard(s.Ctx, &p.GetOrCreateShardRequest{
-		ShardID:          shardID,
+		ShardID:          s.ShardID,
 		InitialShardInfo: shardInfo,
 	})
 	s.NoError(err)
 	s.Equal(shardInfo, resp.ShardInfo)
 
-	updateShardInfo := RandomShardInfo(shardID, rangeID+1)
+	updateShardInfo := RandomShardInfo(s.ShardID, rangeID+1)
 	err = s.ShardManager.UpdateShard(s.Ctx, &p.UpdateShardRequest{
 		ShardInfo:       updateShardInfo,
 		PreviousRangeID: rangeID,
 	})
 	s.NoError(err)
 
 	resp, err = s.ShardManager.GetOrCreateShard(s.Ctx, &p.GetOrCreateShardRequest{
-		ShardID:          shardID,
+		ShardID:          s.ShardID,
 		InitialShardInfo: shardInfo,
 	})
 	s.NoError(err)