use concrete time types

cosmos · Nov 3, 2023 · 0b4898c · 0b4898c
1 parent e818975
commit 0b4898c
Show file tree

Hide file tree

Showing 4 changed files with 71 additions and 49 deletions.
diff --git a/x/protocolpool/keeper/grpc_query_test.go b/x/protocolpool/keeper/grpc_query_test.go
@@ -1,6 +1,8 @@
 package keeper_test
 
 import (
+	"time"
+
 	"cosmossdk.io/x/protocolpool/keeper"
 	"cosmossdk.io/x/protocolpool/types"
 
@@ -9,6 +11,8 @@ import (
 
 func (suite *KeeperTestSuite) TestUnclaimedBudget() {
 	queryServer := keeper.NewQuerier(suite.poolKeeper)
+	startTime := suite.ctx.BlockTime().Add(-70 * time.Second)
+	period := time.Duration(60) * time.Second
 	testCases := []struct {
 		name           string
 		preRun         func()
@@ -40,9 +44,9 @@ func (suite *KeeperTestSuite) TestUnclaimedBudget() {
 				budget := types.Budget{
 					RecipientAddress: recipientAddr.String(),
 					TotalBudget:      &fooCoin,
-					StartTime:        uint64(suite.ctx.BlockTime().Unix() - 70),
+					StartTime:        &startTime,
 					Tranches:         2,
-					Period:           60,
+					Period:           &period,
 				}
 				err := suite.poolKeeper.BudgetProposal.Set(suite.ctx, recipientAddr, budget)
 				suite.Require().NoError(err)
@@ -60,9 +64,9 @@ func (suite *KeeperTestSuite) TestUnclaimedBudget() {
 				budget := types.Budget{
 					RecipientAddress: recipientAddr.String(),
 					TotalBudget:      &fooCoin,
-					StartTime:        uint64(suite.ctx.BlockTime().Unix() - 70),
+					StartTime:        &startTime,
 					Tranches:         2,
-					Period:           60,
+					Period:           &period,
 				}
 				err := suite.poolKeeper.BudgetProposal.Set(suite.ctx, recipientAddr, budget)
 				suite.Require().NoError(err)

diff --git a/x/protocolpool/keeper/keeper.go b/x/protocolpool/keeper/keeper.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"time"
 
 	"cosmossdk.io/collections"
 	storetypes "cosmossdk.io/core/store"
@@ -129,15 +130,15 @@ func (k Keeper) getClaimableFunds(ctx context.Context, recipient sdk.AccAddress)
 		}
 	}
 
-	currentTime := sdkCtx.BlockTime().Unix()
+	currentTime := sdkCtx.BlockTime()
 	startTime := budget.StartTime
 
 	// Check if the start time is reached
-	if uint64(currentTime) < startTime {
+	if currentTime.Before(*startTime) {
 		return sdk.Coin{}, fmt.Errorf("distribution has not started yet")
 	}
 
-	if budget.NextClaimFrom == 0 {
+	if budget.NextClaimFrom == nil || budget.NextClaimFrom.IsZero() {
 		budget.NextClaimFrom = budget.StartTime
 	}
 
@@ -150,31 +151,32 @@ func (k Keeper) getClaimableFunds(ctx context.Context, recipient sdk.AccAddress)
 	return k.calculateClaimableFunds(ctx, recipient, budget, currentTime)
 }
 
-func (k Keeper) calculateClaimableFunds(ctx context.Context, recipient sdk.AccAddress, budget types.Budget, currentTime int64) (amount sdk.Coin, err error) {
+func (k Keeper) calculateClaimableFunds(ctx context.Context, recipient sdk.AccAddress, budget types.Budget, currentTime time.Time) (amount sdk.Coin, err error) {
 	// Calculate the time elapsed since the last claim time
-	timeElapsed := uint64(currentTime) - budget.NextClaimFrom
+	timeElapsed := currentTime.Sub(*budget.NextClaimFrom)
 
 	// Check the time elapsed has passed period length
-	if timeElapsed < budget.Period {
+	if timeElapsed < *budget.Period {
 		return sdk.Coin{}, fmt.Errorf("budget period has not passed yet")
 	}
 
 	// Calculate how many periods have passed
-	periodsPassed := timeElapsed / budget.Period
+	periodsPassed := int64(timeElapsed) / int64(*budget.Period)
 
 	// Calculate the amount to distribute for all passed periods
-	coinsToDistribute := math.NewInt(int64(periodsPassed)).Mul(budget.TotalBudget.Amount.QuoRaw(int64(budget.Tranches)))
+	coinsToDistribute := math.NewInt(periodsPassed).Mul(budget.TotalBudget.Amount.QuoRaw(int64(budget.Tranches)))
 	amount = sdk.NewCoin(budget.TotalBudget.Denom, coinsToDistribute)
 
 	// update the budget's remaining tranches
-	budget.TranchesLeft -= periodsPassed
+	budget.TranchesLeft -= uint64(periodsPassed)
 
 	// update the ClaimedAmount
 	claimedAmount := budget.ClaimedAmount.Add(amount)
 	budget.ClaimedAmount = &claimedAmount
 
 	// Update the last claim time for the budget
-	budget.NextClaimFrom += budget.Period
+	nextClaimFrom := budget.NextClaimFrom.Add(*budget.Period)
+	budget.NextClaimFrom = &nextClaimFrom
 
 	k.Logger(ctx).Debug(fmt.Sprintf("Processing budget for recipient: %s. Amount: %s", budget.RecipientAddress, coinsToDistribute.String()))
 
@@ -195,20 +197,21 @@ func (k Keeper) validateAndUpdateBudgetProposal(ctx context.Context, bp types.Ms
 		return fmt.Errorf("invalid budget proposal: %w", err)
 	}
 
-	currentTime := sdk.UnwrapSDKContext(ctx).BlockTime().Unix()
-	if int64(bp.StartTime) == 0 {
-		bp.StartTime = uint64(currentTime)
+	currentTime := sdk.UnwrapSDKContext(ctx).BlockTime()
+	if bp.StartTime.IsZero() || bp.StartTime == nil {
+		bp.StartTime = &currentTime
 	}
 
-	if bp.StartTime < uint64(currentTime) {
+	// if bp.StartTime < uint64(currentTime) {
+	if currentTime.After(*bp.StartTime) {
 		return fmt.Errorf("invalid budget proposal: start time cannot be less than the current block time")
 	}
 
 	if bp.Tranches == 0 {
 		return fmt.Errorf("invalid budget proposal: tranches must be greater than zero")
 	}
 
-	if bp.Period == 0 {
+	if bp.Period == nil || *bp.Period == 0 {
 		return fmt.Errorf("invalid budget proposal: period length should be greater than zero")
 	}
 

diff --git a/x/protocolpool/keeper/msg_server_test.go b/x/protocolpool/keeper/msg_server_test.go
@@ -18,6 +18,10 @@ var (
 
 func (suite *KeeperTestSuite) TestMsgSubmitBudgetProposal() {
 	invalidCoin := sdk.NewInt64Coin("foo", 0)
+	startTime := suite.ctx.BlockTime().Add(10 * time.Second)
+	invalidStartTime := suite.ctx.BlockTime().Add(-15 * time.Second)
+	period := time.Duration(60) * time.Second
+	zeroPeriod := time.Duration(0) * time.Second
 	testCases := map[string]struct {
 		preRun    func()
 		input     *types.MsgSubmitBudgetProposal
@@ -29,9 +33,9 @@ func (suite *KeeperTestSuite) TestMsgSubmitBudgetProposal() {
 				Authority:        suite.poolKeeper.GetAuthority(),
 				RecipientAddress: "",
 				TotalBudget:      &fooCoin,
-				StartTime:        uint64(suite.ctx.BlockTime().Unix()),
+				StartTime:        &startTime,
 				Tranches:         2,
-				Period:           60,
+				Period:           &period,
 			},
 			expErr:    true,
 			expErrMsg: "empty address string is not allowed",
@@ -41,9 +45,9 @@ func (suite *KeeperTestSuite) TestMsgSubmitBudgetProposal() {
 				Authority:        "",
 				RecipientAddress: recipientAddr.String(),
 				TotalBudget:      &fooCoin,
-				StartTime:        uint64(suite.ctx.BlockTime().Unix()),
+				StartTime:        &startTime,
 				Tranches:         2,
-				Period:           60,
+				Period:           &period,
 			},
 			expErr:    true,
 			expErrMsg: "empty address string is not allowed",
@@ -53,9 +57,9 @@ func (suite *KeeperTestSuite) TestMsgSubmitBudgetProposal() {
 				Authority:        "invalid_authority",
 				RecipientAddress: recipientAddr.String(),
 				TotalBudget:      &fooCoin,
-				StartTime:        uint64(suite.ctx.BlockTime().Unix()),
+				StartTime:        &startTime,
 				Tranches:         2,
-				Period:           60,
+				Period:           &period,
 			},
 			expErr:    true,
 			expErrMsg: "invalid authority",
@@ -65,9 +69,9 @@ func (suite *KeeperTestSuite) TestMsgSubmitBudgetProposal() {
 				Authority:        suite.poolKeeper.GetAuthority(),
 				RecipientAddress: recipientAddr.String(),
 				TotalBudget:      &invalidCoin,
-				StartTime:        uint64(suite.ctx.BlockTime().Unix()),
+				StartTime:        &startTime,
 				Tranches:         2,
-				Period:           60,
+				Period:           &period,
 			},
 			expErr:    true,
 			expErrMsg: "total budget cannot be zero",
@@ -77,9 +81,9 @@ func (suite *KeeperTestSuite) TestMsgSubmitBudgetProposal() {
 				Authority:        suite.poolKeeper.GetAuthority(),
 				RecipientAddress: recipientAddr.String(),
 				TotalBudget:      &fooCoin,
-				StartTime:        uint64(suite.ctx.BlockTime().Unix() - 5),
+				StartTime:        &invalidStartTime,
 				Tranches:         2,
-				Period:           60,
+				Period:           &period,
 			},
 			expErr:    true,
 			expErrMsg: "start time cannot be less than the current block time",
@@ -89,9 +93,9 @@ func (suite *KeeperTestSuite) TestMsgSubmitBudgetProposal() {
 				Authority:        suite.poolKeeper.GetAuthority(),
 				RecipientAddress: recipientAddr.String(),
 				TotalBudget:      &fooCoin,
-				StartTime:        uint64(suite.ctx.BlockTime().Unix() + 10),
+				StartTime:        &startTime,
 				Tranches:         0,
-				Period:           60,
+				Period:           &period,
 			},
 			expErr:    true,
 			expErrMsg: "tranches must be greater than zero",
@@ -101,9 +105,9 @@ func (suite *KeeperTestSuite) TestMsgSubmitBudgetProposal() {
 				Authority:        suite.poolKeeper.GetAuthority(),
 				RecipientAddress: recipientAddr.String(),
 				TotalBudget:      &fooCoin,
-				StartTime:        uint64(suite.ctx.BlockTime().Unix() + 10),
+				StartTime:        &startTime,
 				Tranches:         2,
-				Period:           0,
+				Period:           &zeroPeriod,
 			},
 			expErr:    true,
 			expErrMsg: "period length should be greater than zero",
@@ -113,16 +117,17 @@ func (suite *KeeperTestSuite) TestMsgSubmitBudgetProposal() {
 				Authority:        suite.poolKeeper.GetAuthority(),
 				RecipientAddress: recipientAddr.String(),
 				TotalBudget:      &fooCoin,
-				StartTime:        uint64(suite.ctx.BlockTime().Unix() + 10),
+				StartTime:        &startTime,
 				Tranches:         2,
-				Period:           60,
+				Period:           &period,
 			},
 			expErr: false,
 		},
 	}
 
 	for name, tc := range testCases {
 		suite.Run(name, func() {
+			suite.SetupTest()
 			if tc.preRun != nil {
 				tc.preRun()
 			}
@@ -138,6 +143,9 @@ func (suite *KeeperTestSuite) TestMsgSubmitBudgetProposal() {
 }
 
 func (suite *KeeperTestSuite) TestMsgClaimBudget() {
+	startTime := suite.ctx.BlockTime().Add(-70 * time.Second)
+	period := time.Duration(60) * time.Second
+
 	testCases := map[string]struct {
 		preRun           func()
 		recipientAddress sdk.AccAddress
@@ -157,13 +165,14 @@ func (suite *KeeperTestSuite) TestMsgClaimBudget() {
 		},
 		"claiming before start time": {
 			preRun: func() {
+				startTime := suite.ctx.BlockTime().Add(3600 * time.Second)
 				// Prepare the budget proposal with a future start time
 				budget := types.Budget{
 					RecipientAddress: recipientAddr.String(),
 					TotalBudget:      &fooCoin,
-					StartTime:        uint64(suite.ctx.BlockTime().Unix() + 3600),
+					StartTime:        &startTime,
 					Tranches:         2,
-					Period:           60,
+					Period:           &period,
 				}
 				err := suite.poolKeeper.BudgetProposal.Set(suite.ctx, recipientAddr, budget)
 				suite.Require().NoError(err)
@@ -174,13 +183,14 @@ func (suite *KeeperTestSuite) TestMsgClaimBudget() {
 		},
 		"budget period has not passed": {
 			preRun: func() {
+				startTime := suite.ctx.BlockTime().Add(-50 * time.Second)
 				// Prepare the budget proposal with start time and a short period
 				budget := types.Budget{
 					RecipientAddress: recipientAddr.String(),
 					TotalBudget:      &fooCoin,
-					StartTime:        uint64(suite.ctx.BlockTime().Unix() - 50),
+					StartTime:        &startTime,
 					Tranches:         1,
-					Period:           60,
+					Period:           &period,
 				}
 				err := suite.poolKeeper.BudgetProposal.Set(suite.ctx, recipientAddr, budget)
 				suite.Require().NoError(err)
@@ -195,9 +205,9 @@ func (suite *KeeperTestSuite) TestMsgClaimBudget() {
 				budget := types.Budget{
 					RecipientAddress: recipientAddr.String(),
 					TotalBudget:      &fooCoin,
-					StartTime:        uint64(suite.ctx.BlockTime().Unix() - 70),
+					StartTime:        &startTime,
 					Tranches:         2,
-					Period:           60,
+					Period:           &period,
 				}
 				err := suite.poolKeeper.BudgetProposal.Set(suite.ctx, recipientAddr, budget)
 				suite.Require().NoError(err)
@@ -212,9 +222,9 @@ func (suite *KeeperTestSuite) TestMsgClaimBudget() {
 				budget := types.Budget{
 					RecipientAddress: recipientAddr.String(),
 					TotalBudget:      &fooCoin,
-					StartTime:        uint64(suite.ctx.BlockTime().Unix() - 70),
+					StartTime:        &startTime,
 					Tranches:         2,
-					Period:           60,
+					Period:           &period,
 				}
 				err := suite.poolKeeper.BudgetProposal.Set(suite.ctx, recipientAddr, budget)
 				suite.Require().NoError(err)
@@ -234,13 +244,15 @@ func (suite *KeeperTestSuite) TestMsgClaimBudget() {
 		"valid double claim attempt": {
 			preRun: func() {
 				oneMonthInSeconds := int64(30 * 24 * 60 * 60) // Approximate number of seconds in 1 month
+				startTimeBeforeMonth := suite.ctx.BlockTime().Add(time.Duration(-oneMonthInSeconds) * time.Second)
+				oneMonthPeriod := time.Duration(oneMonthInSeconds) * time.Second
 				// Prepare the budget proposal with valid start time and period of 1 month (in seconds)
 				budget := types.Budget{
 					RecipientAddress: recipientAddr.String(),
 					TotalBudget:      &fooCoin,
-					StartTime:        uint64(suite.ctx.BlockTime().Unix() - oneMonthInSeconds),
+					StartTime:        &startTimeBeforeMonth,
 					Tranches:         2,
-					Period:           uint64(oneMonthInSeconds),
+					Period:           &oneMonthPeriod,
 				}
 				err := suite.poolKeeper.BudgetProposal.Set(suite.ctx, recipientAddr, budget)
 				suite.Require().NoError(err)
@@ -269,9 +281,9 @@ func (suite *KeeperTestSuite) TestMsgClaimBudget() {
 				budget := types.Budget{
 					RecipientAddress: recipientAddr.String(),
 					TotalBudget:      &fooCoin,
-					StartTime:        uint64(suite.ctx.BlockTime().Unix() - 70),
+					StartTime:        &startTime,
 					Tranches:         2,
-					Period:           60,
+					Period:           &period,
 				}
 				err := suite.poolKeeper.BudgetProposal.Set(suite.ctx, recipientAddr, budget)
 				suite.Require().NoError(err)
@@ -290,7 +302,7 @@ func (suite *KeeperTestSuite) TestMsgClaimBudget() {
 					Time: newBlockTime,
 				})
 
-				// Claim the funds twitce
+				// Claim the funds twice
 				msg = &types.MsgClaimBudget{
 					RecipientAddress: recipientAddr.String(),
 				}
@@ -306,6 +318,7 @@ func (suite *KeeperTestSuite) TestMsgClaimBudget() {
 
 	for name, tc := range testCases {
 		suite.Run(name, func() {
+			suite.SetupTest()
 			if tc.preRun != nil {
 				tc.preRun()
 			}

diff --git a/x/protocolpool/types/msg.go b/x/protocolpool/types/msg.go
@@ -1,6 +1,8 @@
 package types
 
 import (
+	"time"
+
 	sdk "github.com/cosmos/cosmos-sdk/types"
 )
 
@@ -28,7 +30,7 @@ func NewCommunityPoolSpend(amount sdk.Coins, authority, recipient string) *MsgCo
 	}
 }
 
-func NewBudgetProposal(recipient string, totalBudget sdk.Coin, startTime, tranches, period uint64) *Budget {
+func NewBudgetProposal(recipient string, totalBudget sdk.Coin, startTime *time.Time, tranches uint64, period *time.Duration) *Budget {
 	return &Budget{
 		RecipientAddress: recipient,
 		TotalBudget:      &totalBudget,