add more unit tests

x/storage/keeper/bucket_rate_limit.go

@@ -19,13 +19,8 @@ func (k Keeper) SetBucketFlowRateLimit(ctx sdk.Context, operator sdk.AccAddress,
 	// get the bucket
 	bucket, found := k.GetBucketInfo(ctx, bucketName)
 
-	// check the bucket owner is the same as the bucket owner of the bucket
-	if found && bucket.Owner != bucketOwner.String() {
-		return fmt.Errorf("the bucket owner is not the same as the bucket owner of the bucket")
-	}
-
 	// if the bucket does not use the payment account, just set the flow rate limit
-	if !found || (found && bucket.PaymentAddress != paymentAccount.String()) {
+	if !found || (found && (bucket.PaymentAddress != paymentAccount.String() || bucket.Owner != bucketOwner.String())) {
 		// set the flow rate limit to the store
 		k.setBucketFlowRateLimit(ctx, paymentAccount, bucketOwner, bucketName, &types.BucketFlowRateLimit{
 			FlowRateLimit: rateLimit,
@@ -107,7 +102,7 @@ func (k Keeper) setFlowRateLimit(ctx sdk.Context, bucketInfo *types.BucketInfo,
 	}
 
 	internalBucketInfo := k.MustGetInternalBucketInfo(ctx, bucketInfo.Id)
-	isRateLimited := k.isBucketRateLimited(ctx, bucketName)
+	isRateLimited := k.IsBucketRateLimited(ctx, bucketName)
 
 	if found && isRateLimited {
 		internalBucketInfo.PriceTime = ctx.BlockTime().Unix()
@@ -211,8 +206,8 @@ func (k Keeper) deleteBucketFlowRateLimitStatus(ctx sdk.Context, bucketName stri
 	}
 }
 
-// isBucketRateLimited checks if the bucket is rate limited
-func (k Keeper) isBucketRateLimited(ctx sdk.Context, bucketName string) bool {
+// IsBucketRateLimited checks if the bucket is rate limited
+func (k Keeper) IsBucketRateLimited(ctx sdk.Context, bucketName string) bool {
 	status, found := k.getBucketFlowRateLimitStatus(ctx, bucketName)
 	if !found {
 		return false

x/storage/keeper/bucket_rate_limit_test.go

@@ -27,7 +27,7 @@ func (s *TestSuite) TestSetBucketFlowRateLimit() {
 	// case 2: bucket is not found
 	s.paymentKeeper.EXPECT().IsPaymentAccountOwner(gomock.Any(), gomock.Any(), gomock.Any()).Return(true).AnyTimes()
 	err = s.storageKeeper.SetBucketFlowRateLimit(s.ctx, operatorAddress, bucketOwner, paymentAccount, bucketName, sdkmath.NewInt(1))
-	s.Require().ErrorContains(err, "No such bucket")
+	s.Require().NoError(err)
 
 	bucketInfo := &types.BucketInfo{
 		Owner:            bucketOwner.String(),
@@ -41,7 +41,7 @@ func (s *TestSuite) TestSetBucketFlowRateLimit() {
 
 	// case 3: different bucket owner
 	err = s.storageKeeper.SetBucketFlowRateLimit(s.ctx, operatorAddress, sample.RandAccAddress(), paymentAccount, bucketName, sdkmath.NewInt(1))
-	s.Require().ErrorContains(err, "invalid bucket owner")
+	s.Require().NoError(err)
 
 	// case 4: bucket does not use the payment account
 	err = s.storageKeeper.SetBucketFlowRateLimit(s.ctx, operatorAddress, bucketOwner, sample.RandAccAddress(), bucketName, sdkmath.NewInt(1))
@@ -69,7 +69,7 @@ func (s *TestSuite) TestSetZeroBucketFlowRateLimit() {
 	s.Require().NoError(err)
 }
 
-func (s *TestSuite) TestSetNonZeroBucketFlowRateLimit() {
+func (s *TestSuite) TestSetFlowRateLimit_NotLimited() {
 	operatorAddress := sample.RandAccAddress()
 	bucketOwner := sample.RandAccAddress()
 	paymentAccount := sample.RandAccAddress()
@@ -92,31 +92,79 @@ func (s *TestSuite) TestSetNonZeroBucketFlowRateLimit() {
 	s.Require().NoError(err)
 	totalOutFlowRate := getTotalOutFlowRate(bill.Flows)
 
-	// case 1: rate limit does not exist before and is less than total out flow rate
-	err = s.storageKeeper.SetBucketFlowRateLimit(s.ctx, operatorAddress, bucketOwner, paymentAccount, bucketName, totalOutFlowRate.Sub(sdkmath.NewInt(1)))
-	s.Require().ErrorContains(err, "greater than the new rate limit")
-
-	// case 2: rate limit does not exist before and is equal to total out flow rate
-	err = s.storageKeeper.SetBucketFlowRateLimit(s.ctx, operatorAddress, bucketOwner, paymentAccount, bucketName, totalOutFlowRate)
+	// case 1: rate limit does not exist before and is larger than total out flow rate
+	err = s.storageKeeper.SetBucketFlowRateLimit(s.ctx, operatorAddress, bucketOwner, paymentAccount, bucketName, totalOutFlowRate.Add(sdkmath.NewInt(1)))
 	s.Require().NoError(err)
 
-	// case 3: rate limit does not exist before and is greater than total out flow rate
+	isRateLimited := s.storageKeeper.IsBucketRateLimited(s.ctx, bucketName)
+	s.Require().False(isRateLimited)
+
+	// case 2: rate limit exists before and is equal to the previous flow rate limit
 	err = s.storageKeeper.SetBucketFlowRateLimit(s.ctx, operatorAddress, bucketOwner, paymentAccount, bucketName, totalOutFlowRate.Add(sdkmath.NewInt(1)))
 	s.Require().NoError(err)
 
-	// case 4: rate limit exists before and is equal than total out flow rate
-	err = s.storageKeeper.SetBucketFlowRateLimit(s.ctx, operatorAddress, bucketOwner, paymentAccount, bucketName, totalOutFlowRate)
+	isRateLimited = s.storageKeeper.IsBucketRateLimited(s.ctx, bucketName)
+	s.Require().False(isRateLimited)
+
+	// case 3: rate limit exists before and larger than the previous flow rate limit
+	err = s.storageKeeper.SetBucketFlowRateLimit(s.ctx, operatorAddress, bucketOwner, paymentAccount, bucketName, totalOutFlowRate.Add(sdkmath.NewInt(2)))
 	s.Require().NoError(err)
 
-	// case 5: rate limit exists before and is 0, and the new rate limit is less than total out flow rate
-	err = s.storageKeeper.SetBucketFlowRateLimit(s.ctx, operatorAddress, bucketOwner, paymentAccount, bucketName, sdk.NewInt(0))
+	isRateLimited = s.storageKeeper.IsBucketRateLimited(s.ctx, bucketName)
+	s.Require().False(isRateLimited)
+}
+
+func (s *TestSuite) TestSetBucketFlowRateLimit_Limited() {
+	operatorAddress := sample.RandAccAddress()
+	bucketOwner := sample.RandAccAddress()
+	paymentAccount := sample.RandAccAddress()
+	bucketName := string(sample.RandStr(10))
+
+	bucketInfo := &types.BucketInfo{
+		Owner:            bucketOwner.String(),
+		BucketName:       bucketName,
+		Id:               sdk.NewUint(1),
+		PaymentAddress:   paymentAccount.String(),
+		ChargedReadQuota: 100,
+	}
+	prepareReadStoreBill(s, bucketInfo)
+
+	s.paymentKeeper.EXPECT().IsPaymentAccountOwner(gomock.Any(), gomock.Any(), gomock.Any()).Return(true).AnyTimes()
+	s.paymentKeeper.EXPECT().ApplyUserFlowsList(gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
+
+	internalBucketInfo := s.storageKeeper.MustGetInternalBucketInfo(s.ctx, bucketInfo.Id)
+	bill, err := s.storageKeeper.GetBucketReadStoreBill(s.ctx, bucketInfo, internalBucketInfo)
+	s.Require().NoError(err)
+	totalOutFlowRate := getTotalOutFlowRate(bill.Flows)
+
+	// case 1: rate limit does not exist before and is less than total out flow rate
+	err = s.storageKeeper.SetBucketFlowRateLimit(s.ctx, operatorAddress, bucketOwner, paymentAccount, bucketName, totalOutFlowRate.Sub(sdkmath.NewInt(1)))
 	s.Require().NoError(err)
+
+	isRateLimited := s.storageKeeper.IsBucketRateLimited(s.ctx, bucketName)
+	s.Require().True(isRateLimited)
+
+	// case 2: rate limit exists before and is equal to the previous flow rate limit
 	err = s.storageKeeper.SetBucketFlowRateLimit(s.ctx, operatorAddress, bucketOwner, paymentAccount, bucketName, totalOutFlowRate.Sub(sdkmath.NewInt(1)))
-	s.Require().ErrorContains(err, "greater than the new rate limit")
+	s.Require().NoError(err)
+
+	isRateLimited = s.storageKeeper.IsBucketRateLimited(s.ctx, bucketName)
+	s.Require().True(isRateLimited)
+
+	// case 3: bucket is rate limited and the new rate limit is less than the total out flow rate
+	err = s.storageKeeper.SetBucketFlowRateLimit(s.ctx, operatorAddress, bucketOwner, paymentAccount, bucketName, totalOutFlowRate.Sub(sdkmath.NewInt(2)))
+	s.Require().NoError(err)
 
-	// case 6: rate limit exists before and is 0, and the new rate limit is larger than the total out flow rate
+	isRateLimited = s.storageKeeper.IsBucketRateLimited(s.ctx, bucketName)
+	s.Require().True(isRateLimited)
+
+	// case 4: bucket is rate limited and the new rate limit is larger than the total out flow rate
 	err = s.storageKeeper.SetBucketFlowRateLimit(s.ctx, operatorAddress, bucketOwner, paymentAccount, bucketName, totalOutFlowRate.Add(sdkmath.NewInt(1)))
 	s.Require().NoError(err)
+
+	// the bucket should be no longer rate limited
+	isRateLimited = s.storageKeeper.IsBucketRateLimited(s.ctx, bucketName)
+	s.Require().False(isRateLimited)
 }
 
 func getTotalOutFlowRate(flows []paymenttypes.OutFlow) sdkmath.Int {
@@ -195,3 +243,7 @@ func prepareReadStoreBill(s *TestSuite, bucketInfo *types.BucketInfo) {
 	s.storageKeeper.StoreBucketInfo(s.ctx, bucketInfo)
 	s.storageKeeper.SetInternalBucketInfo(s.ctx, bucketInfo.Id, internalBucketInfo)
 }
+
+func (s *TestSuite) TestChargeBucketReadFee() {
+
+}

x/storage/keeper/payment.go

@@ -48,7 +48,7 @@ func (k Keeper) UnChargeBucketReadFee(ctx sdk.Context, bucketInfo *storagetypes.
 
 	if ctx.IsUpgraded(upgradetypes.Serengeti) {
 		// if the bucket's flow rate limit is set to zero, no need to uncharge, since the bucket is already uncharged
-		if k.isBucketRateLimited(ctx, bucketInfo.BucketName) {
+		if k.IsBucketRateLimited(ctx, bucketInfo.BucketName) {
 			return nil
 		}
 	}
@@ -347,7 +347,7 @@ func (k Keeper) ChargeViaBucketChange(ctx sdk.Context, bucketInfo *storagetypes.
 	}
 
 	if ctx.IsUpgraded(upgradetypes.Serengeti) {
-		isPreviousBucketLimited := k.isBucketRateLimited(ctx, bucketInfo.BucketName)
+		isPreviousBucketLimited := k.IsBucketRateLimited(ctx, bucketInfo.BucketName)
 
 		if prevPaymentAccount == bucketInfo.PaymentAddress && isPreviousBucketLimited {
 			return fmt.Errorf("payment account is not changed but the bucket is limited")
@@ -550,7 +550,7 @@ func (k Keeper) UnChargeBucketReadStoreFee(ctx sdk.Context, bucketInfo *storaget
 
 	if ctx.IsUpgraded(upgradetypes.Serengeti) {
 		// if the bucket's flow rate limit is set to zero, no need to uncharge, since the bucket is already uncharged
-		if k.isBucketRateLimited(ctx, bucketInfo.BucketName) {
+		if k.IsBucketRateLimited(ctx, bucketInfo.BucketName) {
 			return nil
 		}
 	}