add tests for discontinue count

bnb-chain · Jul 31, 2023 · eccaf83 · eccaf83
1 parent 1735dfc
commit eccaf83
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 8 deletions.
diff --git a/x/storage/keeper/keeper.go b/x/storage/keeper/keeper.go
@@ -446,7 +446,7 @@ func (k Keeper) DiscontinueBucket(ctx sdk.Context, operator sdk.AccAddress, buck
 			"only primary sp is allowed to do discontinue bucket, expect sp id : %d", spInState.Id)
 	}
 
-	count := k.getDiscontinueBucketCount(ctx, operator)
+	count := k.GetDiscontinueBucketCount(ctx, operator)
 	max := k.DiscontinueBucketMax(ctx)
 	if count+1 > max {
 		return types.ErrNoMoreDiscontinue.Wrapf("no more buckets can be requested in this window")
@@ -461,7 +461,7 @@ func (k Keeper) DiscontinueBucket(ctx sdk.Context, operator sdk.AccAddress, buck
 	deleteAt := ctx.BlockTime().Unix() + k.DiscontinueConfirmPeriod(ctx)
 
 	k.appendDiscontinueBucketIds(ctx, deleteAt, []sdkmath.Uint{bucketInfo.Id})
-	k.setDiscontinueBucketCount(ctx, operator, count+1)
+	k.SetDiscontinueBucketCount(ctx, operator, count+1)
 
 	if err := ctx.EventManager().EmitTypedEvents(&types.EventDiscontinueBucket{
 		BucketId:   bucketInfo.Id,
@@ -1126,7 +1126,7 @@ func (k Keeper) DiscontinueObject(ctx sdk.Context, operator sdk.AccAddress, buck
 		return errors.Wrapf(types.ErrAccessDenied, "only primary sp is allowed to do discontinue objects")
 	}
 
-	count := k.getDiscontinueObjectCount(ctx, operator)
+	count := k.GetDiscontinueObjectCount(ctx, operator)
 	max := k.DiscontinueObjectMax(ctx)
 	if count+uint64(len(objectIds)) > max {
 		return types.ErrNoMoreDiscontinue.Wrapf("only %d objects can be requested in this window", max-count)
@@ -1155,7 +1155,7 @@ func (k Keeper) DiscontinueObject(ctx sdk.Context, operator sdk.AccAddress, buck
 
 	deleteAt := ctx.BlockTime().Unix() + k.DiscontinueConfirmPeriod(ctx)
 	k.AppendDiscontinueObjectIds(ctx, deleteAt, objectIds)
-	k.setDiscontinueObjectCount(ctx, operator, count+uint64(len(objectIds)))
+	k.SetDiscontinueObjectCount(ctx, operator, count+uint64(len(objectIds)))
 
 	events := make([]proto.Message, 0)
 	for _, objectId := range objectIds {
@@ -1470,7 +1470,7 @@ func (k Keeper) isNonEmptyBucket(ctx sdk.Context, bucketName string) bool {
 	return iter.Valid()
 }
 
-func (k Keeper) getDiscontinueObjectCount(ctx sdk.Context, operator sdk.AccAddress) uint64 {
+func (k Keeper) GetDiscontinueObjectCount(ctx sdk.Context, operator sdk.AccAddress) uint64 {
 	store := prefix.NewStore(ctx.KVStore(k.storeKey), types.DiscontinueObjectCountPrefix)
 	bz := store.Get(operator.Bytes())
 
@@ -1480,7 +1480,7 @@ func (k Keeper) getDiscontinueObjectCount(ctx sdk.Context, operator sdk.AccAddre
 	return binary.BigEndian.Uint64(bz)
 }
 
-func (k Keeper) setDiscontinueObjectCount(ctx sdk.Context, operator sdk.AccAddress, count uint64) {
+func (k Keeper) SetDiscontinueObjectCount(ctx sdk.Context, operator sdk.AccAddress, count uint64) {
 	store := prefix.NewStore(ctx.KVStore(k.storeKey), types.DiscontinueObjectCountPrefix)
 
 	countBytes := make([]byte, 8)
@@ -1551,7 +1551,7 @@ func (k Keeper) DeleteDiscontinueObjectsUntil(ctx sdk.Context, timestamp int64,
 	return deleted, nil
 }
 
-func (k Keeper) getDiscontinueBucketCount(ctx sdk.Context, operator sdk.AccAddress) uint64 {
+func (k Keeper) GetDiscontinueBucketCount(ctx sdk.Context, operator sdk.AccAddress) uint64 {
 	store := prefix.NewStore(ctx.KVStore(k.storeKey), types.DiscontinueBucketCountPrefix)
 	bz := store.Get(operator.Bytes())
 
@@ -1561,7 +1561,7 @@ func (k Keeper) getDiscontinueBucketCount(ctx sdk.Context, operator sdk.AccAddre
 	return binary.BigEndian.Uint64(bz)
 }
 
-func (k Keeper) setDiscontinueBucketCount(ctx sdk.Context, operator sdk.AccAddress, count uint64) {
+func (k Keeper) SetDiscontinueBucketCount(ctx sdk.Context, operator sdk.AccAddress, count uint64) {
 	store := prefix.NewStore(ctx.KVStore(k.storeKey), types.DiscontinueBucketCountPrefix)
 
 	countBytes := make([]byte, 8)

diff --git a/x/storage/keeper/keeper_test.go b/x/storage/keeper/keeper_test.go
@@ -0,0 +1,29 @@
+package keeper_test
+
+import "github.com/bnb-chain/greenfield/testutil/sample"
+
+func (s *TestSuite) TestClearDiscontinueBucketCount() {
+	acc1 := sample.RandAccAddress()
+	s.storageKeeper.SetDiscontinueBucketCount(s.ctx, acc1, 1)
+
+	count := s.storageKeeper.GetDiscontinueBucketCount(s.ctx, acc1)
+	s.Require().Equal(uint64(1), count)
+
+	s.storageKeeper.ClearDiscontinueBucketCount(s.ctx)
+
+	count = s.storageKeeper.GetDiscontinueBucketCount(s.ctx, acc1)
+	s.Require().Equal(uint64(0), count)
+}
+
+func (s *TestSuite) TestClearDiscontinueObjectCount() {
+	acc1 := sample.RandAccAddress()
+	s.storageKeeper.SetDiscontinueObjectCount(s.ctx, acc1, 1)
+
+	count := s.storageKeeper.GetDiscontinueObjectCount(s.ctx, acc1)
+	s.Require().Equal(uint64(1), count)
+
+	s.storageKeeper.ClearDiscontinueObjectCount(s.ctx)
+
+	count = s.storageKeeper.GetDiscontinueObjectCount(s.ctx, acc1)
+	s.Require().Equal(uint64(0), count)
+}