chore: add unit-test

0xPolygon · Jan 13, 2025 · fc739a8 · fc739a8
1 parent 56e9cf1
commit fc739a8
Show file tree

Hide file tree

Showing 2 changed files with 62 additions and 9 deletions.
diff --git a/x/gov/keeper/common_test.go b/x/gov/keeper/common_test.go
@@ -42,6 +42,7 @@ var (
 	_, _, addr   = testdata.KeyTestPubAddr()
 	govAcct      = authtypes.NewModuleAddress(types.ModuleName)
 	distAcct     = authtypes.NewModuleAddress(disttypes.ModuleName)
+	authAcct     = authtypes.NewModuleAddress(authtypes.FeeCollectorName)
 	TestProposal = getTestProposal()
 )
 
@@ -145,26 +146,53 @@ func setupGovKeeper(t *testing.T) (
 }
 
 // trackMockBalances sets up expected calls on the Mock BankKeeper, and also
-// locally tracks accounts balances (not modules balances).
+// locally tracks accounts balances (also modules balances).
 func trackMockBalances(bankKeeper *govtestutil.MockBankKeeper, distributionKeeper *govtestutil.MockDistributionKeeper) {
 	balances := make(map[string]sdk.Coins)
 	balances[distAcct.String()] = sdk.NewCoins(sdk.NewCoin(sdk.DefaultBondDenom, math.NewInt(0)))
 
-	// We don't track module account balances.
-	bankKeeper.EXPECT().MintCoins(gomock.Any(), topupTypes.ModuleName, gomock.Any()).AnyTimes()
-	bankKeeper.EXPECT().BurnCoins(gomock.Any(), types.ModuleName, gomock.Any()).AnyTimes()
-	bankKeeper.EXPECT().SendCoinsFromModuleToModule(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).AnyTimes()
-
-	// But we do track normal account balances.
-	bankKeeper.EXPECT().SendCoinsFromAccountToModule(gomock.Any(), gomock.Any(), types.ModuleName, gomock.Any()).DoAndReturn(func(_ sdk.Context, sender sdk.AccAddress, _ string, coins sdk.Coins) error {
+	bankKeeper.EXPECT().MintCoins(gomock.Any(), topupTypes.ModuleName, gomock.Any()).DoAndReturn(func(_ sdk.Context, moduleName string, coins sdk.Coins) error {
+		balances[moduleName] = balances[moduleName].Add(coins...)
+		return nil
+	}).AnyTimes()
+	bankKeeper.EXPECT().BurnCoins(gomock.Any(), types.ModuleName, gomock.Any()).DoAndReturn(func(_ sdk.Context, moduleName string, coins sdk.Coins) error {
+		newBalance, negative := balances[moduleName].SafeSub(coins...)
+		if negative {
+			return fmt.Errorf("not enough balance")
+		}
+		balances[moduleName] = newBalance
+		return nil
+	}).AnyTimes()
+	bankKeeper.EXPECT().SendCoinsFromModuleToModule(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn(func(_ sdk.Context, senderModule, receiverModule string, coins sdk.Coins) error {
+		// sender module balance
+		newBalance, negative := balances[senderModule].SafeSub(coins...)
+		if negative {
+			return fmt.Errorf("not enough balance")
+		}
+		balances[senderModule] = newBalance
+		// receiver module balance
+		balances[receiverModule] = balances[receiverModule].Add(coins...)
+		return nil
+	}).AnyTimes()
+	bankKeeper.EXPECT().SendCoinsFromAccountToModule(gomock.Any(), gomock.Any(), types.ModuleName, gomock.Any()).DoAndReturn(func(_ sdk.Context, sender sdk.AccAddress, receiverModule string, coins sdk.Coins) error {
+		// sender account balance
 		newBalance, negative := balances[sender.String()].SafeSub(coins...)
 		if negative {
 			return fmt.Errorf("not enough balance")
 		}
 		balances[sender.String()] = newBalance
+		// receiver module balance
+		balances[receiverModule] = balances[receiverModule].Add(coins...)
 		return nil
 	}).AnyTimes()
-	bankKeeper.EXPECT().SendCoinsFromModuleToAccount(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn(func(_ sdk.Context, module string, rcpt sdk.AccAddress, coins sdk.Coins) error {
+	bankKeeper.EXPECT().SendCoinsFromModuleToAccount(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn(func(_ sdk.Context, senderModule string, rcpt sdk.AccAddress, coins sdk.Coins) error {
+		// sender module balance
+		newBalance, negative := balances[senderModule].SafeSub(coins...)
+		if negative {
+			return fmt.Errorf("not enough balance")
+		}
+		balances[senderModule] = newBalance
+		// receiver account balance
 		balances[rcpt.String()] = balances[rcpt.String()].Add(coins...)
 		return nil
 	}).AnyTimes()

diff --git a/x/gov/keeper/deposit_test.go b/x/gov/keeper/deposit_test.go
@@ -157,6 +157,31 @@ func TestDeposits(t *testing.T) {
 			deposits, _ = govKeeper.GetDeposits(ctx, proposalID)
 			require.Len(t, deposits, 1)
 			require.Equal(t, addr0Initial.Sub(fourStake...), bankKeeper.GetAllBalances(ctx, TestAddrs[0]))
+
+			// Test SendToFeeCollectorAndDeleteDeposits
+			addr0Initial = bankKeeper.GetAllBalances(ctx, TestAddrs[0])
+			addr1Initial = bankKeeper.GetAllBalances(ctx, TestAddrs[1])
+			feeCollectorInitial := bankKeeper.GetAllBalances(ctx, authAcct)
+
+			proposal, err = govKeeper.SubmitProposal(ctx, tp, "", "title", "summary", TestAddrs[0], false)
+			require.NoError(t, err)
+			proposalID = proposal.Id
+
+			_, err = govKeeper.AddDeposit(ctx, proposalID, TestAddrs[0], fourStake)
+			require.NoError(t, err)
+
+			_, err = govKeeper.AddDeposit(ctx, proposalID, TestAddrs[1], fiveStake)
+			require.NoError(t, err)
+
+			deposits, _ = govKeeper.GetDeposits(ctx, proposalID)
+			require.Len(t, deposits, 2)
+
+			err = govKeeper.SendToFeeCollectorAndDeleteDeposits(ctx, proposalID)
+			require.NoError(t, err)
+
+			require.Equal(t, addr0Initial.Sub(fourStake...), bankKeeper.GetAllBalances(ctx, TestAddrs[0]))
+			require.Equal(t, addr1Initial.Sub(fiveStake...), bankKeeper.GetAllBalances(ctx, TestAddrs[1]))
+			require.Equal(t, feeCollectorInitial.Add(fourStake...).Add(fiveStake...), bankKeeper.GetAllBalances(ctx, authAcct))
 		})
 	}
 }