fix: fix broken x/foundation invariant on treasury (#946)

* TruncateDecimal first

* Refactor CollectFoundationTax

* Update tests

* Refactor tests

* Update CHANGELOG.md

* Apply feedbacks on the PR

* Revert the changes in CollectFoundationTax's signature

* Ensure the behavior does not change

Author: 0Tech

CHANGELOG.md

@@ -49,6 +49,7 @@ Ref: https://keepachangelog.com/en/1.0.0/
 * (swagger) [\#898](https://github.com/line/lbm-sdk/pull/898) fix a bug not added `lbm.tx.v1beta1.Service/GetBlockWithTxs` in swagger
 * (x/collection) [\#911](https://github.com/line/lbm-sdk/pull/911) Add missing command(TxCmdModify) for CLI
 * (x/foundation) [\#922](https://github.com/line/lbm-sdk/pull/922) Propagate events in x/foundation through sdk.Results
+* (x/foundation) [\#946](https://github.com/line/lbm-sdk/pull/946) Fix broken x/foundation invariant on treasury
 
 ### Removed
 

x/foundation/expected_keepers.go

@@ -17,7 +17,6 @@ type (
 	BankKeeper interface {
 		GetAllBalances(ctx sdk.Context, addr sdk.AccAddress) sdk.Coins
 
-		SendCoinsFromModuleToModule(ctx sdk.Context, senderModule string, recipientModule string, amt sdk.Coins) error
 		SendCoinsFromModuleToAccount(ctx sdk.Context, senderModule string, recipientAddr sdk.AccAddress, amt sdk.Coins) error
 		SendCoinsFromAccountToModule(ctx sdk.Context, senderAddr sdk.AccAddress, recipientModule string, amt sdk.Coins) error
 	}

x/foundation/keeper/internal/abci_test.go

@@ -9,8 +9,9 @@ import (
 func (s *KeeperTestSuite) TestBeginBlocker() {
 	ctx, _ := s.ctx.CacheContext()
 
+	taxRatio := sdk.MustNewDecFromStr("0.123456789")
 	s.impl.SetParams(ctx, foundation.Params{
-		FoundationTax: sdk.MustNewDecFromStr("0.5"),
+		FoundationTax: taxRatio,
 	})
 
 	before := s.impl.GetTreasury(ctx)
@@ -20,10 +21,14 @@ func (s *KeeperTestSuite) TestBeginBlocker() {
 	// collect
 	internal.BeginBlocker(ctx, s.impl)
 
+	tax := sdk.NewDecFromInt(s.balance).MulTruncate(taxRatio).TruncateInt()
+	// ensure the behavior does not change
+	s.Require().Equal(sdk.NewInt(121932631), tax)
+
+	expectedAfter := s.balance.Add(tax)
 	after := s.impl.GetTreasury(ctx)
 	s.Require().Equal(1, len(after))
-	// s.balance + s.balance * 0.5
-	s.Require().Equal(sdk.NewDecFromInt(s.balance.Add(s.balance.Quo(sdk.NewInt(2)))), after[0].Amount)
+	s.Require().Equal(sdk.NewDecFromInt(expectedAfter), after[0].Amount)
 }
 
 func (s *KeeperTestSuite) TestEndBlocker() {

x/foundation/keeper/internal/keeper_test.go

@@ -23,9 +23,10 @@ type KeeperTestSuite struct {
 	suite.Suite
 	ctx sdk.Context
 
-	app             *simapp.SimApp
-	keeper          keeper.Keeper
-	impl            internal.Keeper
+	bankKeeper foundation.BankKeeper
+	keeper     keeper.Keeper
+	impl       internal.Keeper
+
 	queryServer     foundation.QueryServer
 	msgServer       foundation.MsgServer
 	proposalHandler govtypes.Handler
@@ -54,18 +55,19 @@ func newMsgCreateDog(name string) sdk.Msg {
 
 func (s *KeeperTestSuite) SetupTest() {
 	checkTx := false
-	s.app = simapp.Setup(checkTx)
-	testdata.RegisterInterfaces(s.app.InterfaceRegistry())
-	testdata.RegisterMsgServer(s.app.MsgServiceRouter(), testdata.MsgServerImpl{})
+	app := simapp.Setup(checkTx)
+	testdata.RegisterInterfaces(app.InterfaceRegistry())
+	testdata.RegisterMsgServer(app.MsgServiceRouter(), testdata.MsgServerImpl{})
 
-	s.ctx = s.app.BaseApp.NewContext(checkTx, tmproto.Header{})
-	s.keeper = s.app.FoundationKeeper
+	s.ctx = app.BaseApp.NewContext(checkTx, tmproto.Header{})
+	s.bankKeeper = app.BankKeeper
+	s.keeper = app.FoundationKeeper
 	s.impl = internal.NewKeeper(
-		s.app.AppCodec(),
-		s.app.GetKey(foundation.ModuleName),
-		s.app.MsgServiceRouter(),
-		s.app.AccountKeeper,
-		s.app.BankKeeper,
+		app.AppCodec(),
+		app.GetKey(foundation.ModuleName),
+		app.MsgServiceRouter(),
+		app.AccountKeeper,
+		app.BankKeeper,
 		authtypes.FeeCollectorName,
 		foundation.DefaultConfig(),
 		foundation.DefaultAuthority().String(),
@@ -106,14 +108,14 @@ func (s *KeeperTestSuite) SetupTest() {
 	s.Require().NoError(err)
 	s.impl.SetFoundationInfo(s.ctx, info)
 
-	s.balance = sdk.NewInt(1000000)
+	s.balance = sdk.NewInt(987654321)
 	s.impl.SetPool(s.ctx, foundation.Pool{
 		Treasury: sdk.NewDecCoinsFromCoins(sdk.NewCoin(sdk.DefaultBondDenom, s.balance)),
 	})
 	holders := []sdk.AccAddress{
 		s.stranger,
-		s.app.AccountKeeper.GetModuleAccount(s.ctx, foundation.TreasuryName).GetAddress(),
-		s.app.AccountKeeper.GetModuleAccount(s.ctx, authtypes.FeeCollectorName).GetAddress(),
+		app.AccountKeeper.GetModuleAccount(s.ctx, foundation.TreasuryName).GetAddress(),
+		app.AccountKeeper.GetModuleAccount(s.ctx, authtypes.FeeCollectorName).GetAddress(),
 	}
 	for _, holder := range holders {
 		amount := sdk.NewCoins(sdk.NewCoin(sdk.DefaultBondDenom, s.balance))
@@ -123,11 +125,11 @@ func (s *KeeperTestSuite) SetupTest() {
 		// because x/bank already has dependency on x/mint, and we must have dependency
 		// on x/bank, it's OK to use x/mint here.
 		minterName := minttypes.ModuleName
-		err := s.app.BankKeeper.MintCoins(s.ctx, minterName, amount)
+		err := app.BankKeeper.MintCoins(s.ctx, minterName, amount)
 		s.Require().NoError(err)
 
-		minter := s.app.AccountKeeper.GetModuleAccount(s.ctx, minterName).GetAddress()
-		err = s.app.BankKeeper.SendCoins(s.ctx, minter, holder, amount)
+		minter := app.AccountKeeper.GetModuleAccount(s.ctx, minterName).GetAddress()
+		err = app.BankKeeper.SendCoins(s.ctx, minter, holder, amount)
 		s.Require().NoError(err)
 	}
 

x/foundation/keeper/internal/treasury.go

@@ -7,24 +7,16 @@ import (
 )
 
 func (k Keeper) CollectFoundationTax(ctx sdk.Context) error {
-	// fetch and clear the collected fees for the fund, since this is
-	// called in BeginBlock, collected fees will be from the previous block
-	feeCollector := k.authKeeper.GetModuleAccount(ctx, k.feeCollectorName)
-	feesCollectedInt := k.bankKeeper.GetAllBalances(ctx, feeCollector.GetAddress())
+	feeCollector := k.authKeeper.GetModuleAccount(ctx, k.feeCollectorName).GetAddress()
+	feesCollectedInt := k.bankKeeper.GetAllBalances(ctx, feeCollector)
 	feesCollected := sdk.NewDecCoinsFromCoins(feesCollectedInt...)
 
 	// calculate the tax
 	taxRatio := k.GetFoundationTax(ctx)
-	tax := feesCollected.MulDecTruncate(taxRatio)
+	tax, _ := feesCollected.MulDecTruncate(taxRatio).TruncateDecimal()
 
-	// update foundation treasury
-	pool := k.GetPool(ctx)
-	pool.Treasury = pool.Treasury.Add(tax...)
-	k.SetPool(ctx, pool)
-
-	// collect tax to the foundation treasury
-	amount, _ := tax.TruncateDecimal()
-	if err := k.bankKeeper.SendCoinsFromModuleToModule(ctx, k.feeCollectorName, foundation.TreasuryName, amount); err != nil {
+	// collect the tax
+	if err := k.FundTreasury(ctx, feeCollector, tax); err != nil {
 		return err
 	}
 

x/foundation/keeper/internal/treasury_test.go

@@ -2,6 +2,8 @@ package internal_test
 
 import (
 	sdk "github.com/line/lbm-sdk/types"
+	authtypes "github.com/line/lbm-sdk/x/auth/types"
+	"github.com/line/lbm-sdk/x/foundation"
 )
 
 func (s *KeeperTestSuite) TestFundTreasury() {
@@ -32,13 +34,16 @@ func (s *KeeperTestSuite) TestFundTreasury() {
 			}
 			s.Require().NoError(err)
 
-			after := s.impl.GetTreasury(ctx)
-			s.Require().Equal(before.Add(sdk.NewDecCoinsFromCoins(amount...)...), after)
+			expectedAfter := before.Add(sdk.NewDecCoinsFromCoins(amount...)...)
+			poolAfter := s.impl.GetTreasury(ctx)
+			s.Require().Equal(sdk.NewDecCoins(expectedAfter...), sdk.NewDecCoins(poolAfter...))
+			balanceAfter := sdk.NewDecCoinsFromCoins(s.bankKeeper.GetAllBalances(ctx, authtypes.NewModuleAddress(foundation.TreasuryName))...)
+			s.Require().Equal(sdk.NewDecCoins(expectedAfter...), sdk.NewDecCoins(balanceAfter...))
 		})
 	}
 }
 
-func (s *KeeperTestSuite) TestWithDrawFromTreasury() {
+func (s *KeeperTestSuite) TestWithdrawFromTreasury() {
 	testCases := map[string]struct {
 		amount sdk.Int
 		valid  bool
@@ -66,8 +71,11 @@ func (s *KeeperTestSuite) TestWithDrawFromTreasury() {
 			}
 			s.Require().NoError(err)
 
-			after := s.impl.GetTreasury(ctx)
-			s.Require().Equal(before.Sub(sdk.NewDecCoinsFromCoins(amount...)), after)
+			expectedAfter := before.Sub(sdk.NewDecCoinsFromCoins(amount...))
+			poolAfter := s.impl.GetTreasury(ctx)
+			s.Require().Equal(sdk.NewDecCoins(expectedAfter...), sdk.NewDecCoins(poolAfter...))
+			balanceAfter := sdk.NewDecCoinsFromCoins(s.bankKeeper.GetAllBalances(ctx, authtypes.NewModuleAddress(foundation.TreasuryName))...)
+			s.Require().Equal(sdk.NewDecCoins(expectedAfter...), sdk.NewDecCoins(balanceAfter...))
 		})
 	}
 }