From ef01d6c1656a20296fb87c496606dd32444c2758 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?L=E1=BA=B7c?= <expertdicer@gmail.com>
Date: Tue, 1 Aug 2023 12:39:42 +0700
Subject: [PATCH] custom bank total supply logic and test

---
 custom/bank/bank_test.go              | 32 +++++++++++++++++++++++++--
 x/transfermiddleware/keeper/keeper.go |  5 +++--
 2 files changed, 33 insertions(+), 4 deletions(-)

diff --git a/custom/bank/bank_test.go b/custom/bank/bank_test.go
index 252b03ce5..bb72961e4 100644
--- a/custom/bank/bank_test.go
+++ b/custom/bank/bank_test.go
@@ -6,6 +6,7 @@ import (
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	banktypes "github.com/cosmos/cosmos-sdk/x/bank/types"
 	ibctransfertypes "github.com/cosmos/ibc-go/v7/modules/apps/transfer/types"
+	transfertypes "github.com/cosmos/ibc-go/v7/modules/apps/transfer/types"
 	clienttypes "github.com/cosmos/ibc-go/v7/modules/core/02-client/types"
 	customibctesting "github.com/notional-labs/centauri/v4/app/ibctesting"
 	"github.com/stretchr/testify/suite"
@@ -48,24 +49,30 @@ func (suite *CustomBankTestSuite) TestTotalSupply() {
 	var (
 		transferAmount = sdk.NewInt(1000000000)
 		// when transfer via sdk transfer from A (module) -> B (contract)
-		coinToSendToB = sdk.NewCoin(sdk.DefaultBondDenom, transferAmount)
-		timeoutHeight = clienttypes.NewHeight(1, 110)
+		coinToSendToB     = sdk.NewCoin(sdk.DefaultBondDenom, transferAmount)
+		timeoutHeight     = clienttypes.NewHeight(1, 110)
+		originAmt, err    = sdk.NewIntFromString("10000000001100000000000")
+		chainBOriginSuply = sdk.NewCoin("stake", originAmt)
 	)
+	suite.Require().True(err)
 	var (
 		expChainBBalanceDiff sdk.Coin
 		path                 = NewTransferPath(suite.chainA, suite.chainB)
+		escrowAddr           = transfertypes.GetEscrowAddress(transfertypes.PortID, "channel-0")
 	)
 
 	testCases := []struct {
 		name                 string
 		expChainABalanceDiff sdk.Coin
 		expTotalSupplyDiff   sdk.Coins
+		expChainBTotalSuppy  sdk.Coins
 		malleate             func()
 	}{
 		{
 			"Total supply with no transfermiddleware setup",
 			sdk.NewCoin(sdk.DefaultBondDenom, transferAmount),
 			sdk.Coins{sdk.NewCoin("ibc/C053D637CCA2A2BA030E2C5EE1B28A16F71CCB0E45E8BE52766DC1B241B77878", transferAmount)},
+			sdk.Coins{sdk.NewCoin("ibc/C053D637CCA2A2BA030E2C5EE1B28A16F71CCB0E45E8BE52766DC1B241B77878", transferAmount), chainBOriginSuply},
 			func() {
 				expChainBBalanceDiff = ibctransfertypes.GetTransferCoin(path.EndpointB.ChannelConfig.PortID, path.EndpointB.ChannelID, coinToSendToB.Denom, transferAmount)
 			},
@@ -74,6 +81,7 @@ func (suite *CustomBankTestSuite) TestTotalSupply() {
 			"Total supply with transfermiddleware setup",
 			sdk.NewCoin(sdk.DefaultBondDenom, transferAmount),
 			sdk.NewCoins(sdk.NewCoin(sdk.DefaultBondDenom, transferAmount)),
+			sdk.Coins{chainBOriginSuply.Add(sdk.NewCoin("stake", transferAmount))},
 			func() {
 				// Add parachain token info
 				chainBtransMiddleware := suite.chainB.TransferMiddleware()
@@ -82,6 +90,25 @@ func (suite *CustomBankTestSuite) TestTotalSupply() {
 				suite.Require().NoError(err)
 			},
 		},
+		{
+			"Total supply with transfermiddleware setup and pre mint",
+			sdk.NewCoin(sdk.DefaultBondDenom, transferAmount),
+			sdk.NewCoins(sdk.NewCoin(sdk.DefaultBondDenom, transferAmount)),
+			sdk.Coins{chainBOriginSuply.Add(sdk.NewCoin("stake", transferAmount))},
+			func() {
+				// Premint for escrow
+				err := suite.chainB.GetBankKeeper().MintCoins(suite.chainB.GetContext(), "mint", sdk.NewCoins(sdk.NewCoin(sdk.DefaultBondDenom, sdk.NewInt(1000000000))))
+				suite.Require().NoError(err)
+				err = suite.chainB.GetBankKeeper().SendCoinsFromModuleToAccount(suite.chainB.GetContext(), "mint", escrowAddr, sdk.NewCoins(sdk.NewCoin(sdk.DefaultBondDenom, sdk.NewInt(1000000000))))
+				suite.Require().NoError(err)
+
+				// Add parachain token info
+				chainBtransMiddleware := suite.chainB.TransferMiddleware()
+				expChainBBalanceDiff = sdk.NewCoin(sdk.DefaultBondDenom, transferAmount)
+				err = chainBtransMiddleware.AddParachainIBCInfo(suite.chainB.GetContext(), "ibc/C053D637CCA2A2BA030E2C5EE1B28A16F71CCB0E45E8BE52766DC1B241B77878", "channel-0", sdk.DefaultBondDenom, sdk.DefaultBondDenom)
+				suite.Require().NoError(err)
+			},
+		},
 	}
 	for _, tc := range testCases {
 		suite.Run(tc.name, func() {
@@ -128,6 +155,7 @@ func (suite *CustomBankTestSuite) TestTotalSupply() {
 			totalSupply, err := suite.chainB.GetBankKeeper().TotalSupply(suite.chainB.GetContext(), &banktypes.QueryTotalSupplyRequest{})
 			suite.Require().NoError(err)
 			suite.Require().Equal(totalSupply.Supply, originalChainBTotalSupply.Supply.Add(tc.expTotalSupplyDiff...))
+			suite.Require().Equal(totalSupply.Supply, tc.expChainBTotalSuppy)
 		})
 	}
 }
diff --git a/x/transfermiddleware/keeper/keeper.go b/x/transfermiddleware/keeper/keeper.go
index 88423fc5b..dc1bb1b67 100644
--- a/x/transfermiddleware/keeper/keeper.go
+++ b/x/transfermiddleware/keeper/keeper.go
@@ -203,8 +203,9 @@ func (keeper Keeper) GetNativeDenomByIBCDenomSecondaryIndex(ctx sdk.Context, ibc
 
 func (keeper Keeper) GetTotalEscrowedToken(ctx sdk.Context) (coins sdk.Coins) {
 	keeper.IterateParaTokenInfos(ctx, func(index int64, info types.ParachainIBCTokenInfo) (stop bool) {
-		escrowCoin := keeper.bankKeeper.GetBalance(ctx, transfertypes.GetEscrowAddress(transfertypes.PortID, info.ChannelId), info.IbcDenom)
-		coins = append(coins, escrowCoin)
+		escrowIbcCoin := keeper.bankKeeper.GetBalance(ctx, transfertypes.GetEscrowAddress(transfertypes.PortID, info.ChannelId), info.IbcDenom)
+		escrowNativeCoin := keeper.bankKeeper.GetBalance(ctx, transfertypes.GetEscrowAddress(transfertypes.PortID, info.ChannelId), info.NativeDenom)
+		coins = append(coins, escrowIbcCoin, escrowNativeCoin)
 		return false
 	})