elys-network · avkr003 · Dec 4, 2024 · Dec 4, 2024 · Dec 4, 2024 · Dec 4, 2024
diff --git a/x/amm/keeper/abci_test.go b/x/amm/keeper/abci_test.go
@@ -276,6 +276,8 @@ func (suite *AmmKeeperTestSuite) TestExecuteSwapRequests() {
 			}
 			suite.app.AmmKeeper.SetPool(suite.ctx, pool)
 			suite.app.AmmKeeper.SetPool(suite.ctx, pool2)
+			suite.Require().True(suite.VerifyPoolAssetWithBalance(1))
+			suite.Require().True(suite.VerifyPoolAssetWithBalance(2))
 
 			msgServer := keeper.NewMsgServerImpl(*suite.app.AmmKeeper)
 			for _, msg := range tc.swapMsgs {
@@ -299,6 +301,8 @@ func (suite *AmmKeeperTestSuite) TestExecuteSwapRequests() {
 			for i, order := range tc.expSwapOrder {
 				suite.Require().Equal(tc.swapMsgs[order], requests[i], fmt.Sprintf("%dth item not match", i))
 			}
+			suite.Require().True(suite.VerifyPoolAssetWithBalance(1))
+			suite.Require().True(suite.VerifyPoolAssetWithBalance(2))
 		})
 	}
 }

diff --git a/x/amm/keeper/apply_exit_pool_state_change_test.go b/x/amm/keeper/apply_exit_pool_state_change_test.go
@@ -71,6 +71,7 @@ func (suite *AmmKeeperTestSuite) TestApplyExitPoolStateChange() {
 				poolId, err := amm.CreatePool(ctx, msg)
 				suite.Require().NoError(err)
 				suite.Require().Equal(poolId, uint64(1))
+				suite.Require().True(suite.VerifyPoolAssetWithBalance(1))
 
 				pool, found := amm.GetPool(ctx, poolId)
 				suite.Require().True(found)
@@ -157,6 +158,7 @@ func (suite *AmmKeeperTestSuite) TestApplyExitPoolStateChange() {
 				poolId, err := amm.CreatePool(ctx, msg)
 				suite.Require().NoError(err)
 				suite.Require().Equal(poolId, uint64(1))
+				suite.Require().True(suite.VerifyPoolAssetWithBalance(1))
 
 				pool, found := amm.GetPool(ctx, poolId)
 				suite.Require().True(found)

diff --git a/x/amm/keeper/apply_join_pool_state_change_test.go b/x/amm/keeper/apply_join_pool_state_change_test.go
@@ -61,6 +61,7 @@ func (suite *AmmKeeperTestSuite) TestApplyJoinPoolStateChange() {
 				poolId, err := amm.CreatePool(ctx, msg)
 				suite.Require().NoError(err)
 				suite.Require().Equal(poolId, uint64(1))
+				suite.Require().True(suite.VerifyPoolAssetWithBalance(1))
 
 				pool, found := amm.GetPool(ctx, poolId)
 				suite.Require().True(found)
@@ -121,6 +122,7 @@ func (suite *AmmKeeperTestSuite) TestApplyJoinPoolStateChange() {
 				poolId, err := amm.CreatePool(ctx, msg)
 				suite.Require().NoError(err)
 				suite.Require().Equal(poolId, uint64(1))
+				suite.Require().True(suite.VerifyPoolAssetWithBalance(1))
 
 				pool, found := amm.GetPool(ctx, poolId)
 				suite.Require().True(found)
@@ -184,6 +186,7 @@ func (suite *AmmKeeperTestSuite) TestApplyJoinPoolStateChange() {
 				poolId, err := amm.CreatePool(ctx, msg)
 				suite.Require().NoError(err)
 				suite.Require().Equal(poolId, uint64(1))
+				suite.Require().True(suite.VerifyPoolAssetWithBalance(1))
 
 				pool, found := amm.GetPool(ctx, poolId)
 				suite.Require().True(found)
@@ -251,6 +254,7 @@ func (suite *AmmKeeperTestSuite) TestApplyJoinPoolStateChange() {
 				poolId, err := amm.CreatePool(ctx, msg)
 				suite.Require().NoError(err)
 				suite.Require().Equal(poolId, uint64(1))
+				suite.Require().True(suite.VerifyPoolAssetWithBalance(1))
 
 				pool, found := amm.GetPool(ctx, poolId)
 				suite.Require().True(found)

diff --git a/x/amm/keeper/fee_test.go b/x/amm/keeper/fee_test.go
@@ -103,12 +103,15 @@ func (suite *AmmKeeperTestSuite) TestOnCollectFee() {
 				},
 				TotalWeight: sdkmath.ZeroInt(),
 			}
+			suite.app.AmmKeeper.SetPool(suite.ctx, pool)
+			suite.Require().True(suite.VerifyPoolAssetWithBalance(1))
+
 			err = suite.app.AmmKeeper.OnCollectFee(suite.ctx, pool, tc.fee)
 			if !tc.expPass {
 				suite.Require().Error(err)
 			} else {
 				suite.Require().NoError(err)
-
+				suite.Require().True(suite.VerifyPoolAssetWithBalance(1))
 				// check pool balance increase/decrease
 				balances := suite.app.BankKeeper.GetAllBalances(suite.ctx, revenueAddr)
 				suite.Require().Equal(balances.String(), tc.expRevenueBalance.String())
@@ -194,12 +197,14 @@ func (suite *AmmKeeperTestSuite) TestSwapFeesToRevenueToken() {
 				},
 				TotalWeight: sdkmath.ZeroInt(),
 			}
+			suite.app.AmmKeeper.SetPool(suite.ctx, pool)
+			suite.Require().True(suite.VerifyPoolAssetWithBalance(1))
 			err = suite.app.AmmKeeper.SwapFeesToRevenueToken(suite.ctx, pool, tc.fee)
 			if !tc.expPass {
 				suite.Require().Error(err)
 			} else {
 				suite.Require().NoError(err)
-
+				suite.Require().True(suite.VerifyPoolAssetWithBalance(1))
 				// check pool balance increase/decrease
 				balances := suite.app.BankKeeper.GetAllBalances(suite.ctx, revenueAddr)
 				suite.Require().Equal(balances.String(), tc.expRevenueBalance.String())

diff --git a/x/amm/keeper/keeper_swap_exact_amount_in_test.go b/x/amm/keeper/keeper_swap_exact_amount_in_test.go
@@ -240,14 +240,16 @@ func (suite *AmmKeeperTestSuite) TestSwapExactAmountIn() {
 				},
 				TotalWeight: sdkmath.ZeroInt(),
 			}
+			suite.app.AmmKeeper.SetPool(suite.ctx, pool)
+			suite.Require().True(suite.VerifyPoolAssetWithBalance(1))
 
 			tokenOut, err := suite.app.AmmKeeper.InternalSwapExactAmountIn(suite.ctx, sender, recipient, pool, tc.tokenIn, tc.tokenOut.Denom, tc.tokenOutMin, tc.swapFeeIn)
 			if !tc.expPass {
 				suite.Require().Error(err)
 			} else {
 				suite.Require().NoError(err)
 				suite.Require().Equal(tokenOut.String(), tc.tokenOut.Amount.String())
-
+				suite.Require().True(suite.VerifyPoolAssetWithBalance(1))
 				// check pool balance increase/decrease
 				balances := suite.app.BankKeeper.GetAllBalances(suite.ctx, poolAddr)
 				suite.Require().Equal(balances.String(), tc.expPoolBalance.String())

diff --git a/x/amm/keeper/keeper_swap_exact_amount_out_test.go b/x/amm/keeper/keeper_swap_exact_amount_out_test.go
@@ -231,14 +231,16 @@ func (suite *AmmKeeperTestSuite) TestSwapExactAmountOut() {
 				},
 				TotalWeight: sdkmath.ZeroInt(),
 			}
+			suite.app.AmmKeeper.SetPool(suite.ctx, pool)
+			suite.Require().True(suite.VerifyPoolAssetWithBalance(1))
 
 			tokenInAmount, err := suite.app.AmmKeeper.InternalSwapExactAmountOut(suite.ctx, sender, recipient, pool, tc.tokenIn.Denom, tc.tokenInMax, tc.tokenOut, tc.swapFeeOut)
 			if !tc.expPass {
 				suite.Require().Error(err)
 			} else {
 				suite.Require().NoError(err)
 				suite.Require().Equal(tokenInAmount.String(), tc.tokenIn.Amount.String())
-
+				suite.Require().True(suite.VerifyPoolAssetWithBalance(1))
 				// check pool balance increase/decrease
 				balances := suite.app.BankKeeper.GetAllBalances(suite.ctx, poolAddr)
 				suite.Require().Equal(balances.String(), tc.expPoolBalance.String())

diff --git a/x/amm/keeper/keeper_test.go b/x/amm/keeper/keeper_test.go
@@ -1,12 +1,13 @@
 package keeper_test
 
 import (
-	authtypes "github.com/cosmos/cosmos-sdk/x/auth/types"
-	govtypes "github.com/cosmos/cosmos-sdk/x/gov/types"
 	"sort"
 	"strings"
 	"testing"
 
+	authtypes "github.com/cosmos/cosmos-sdk/x/auth/types"
+	govtypes "github.com/cosmos/cosmos-sdk/x/gov/types"
+
 	"cosmossdk.io/math"
 
 	"github.com/cometbft/cometbft/crypto/ed25519"
@@ -299,3 +300,20 @@ func SetupMockPools(k *keeper.Keeper, ctx sdk.Context) {
 		k.SetPool(ctx, pool)
 	}
 }
+
+func (suite *AmmKeeperTestSuite) VerifyPoolAssetWithBalance(poolId uint64) bool {
+	pool, found := suite.app.AmmKeeper.GetPool(suite.ctx, poolId)
+	if !found {
+		return false
+	}
+	for _, asset := range pool.PoolAssets {
+		bal := suite.app.BankKeeper.GetBalance(suite.ctx, sdk.MustAccAddressFromBech32(pool.Address), asset.Token.Denom)
+		if !asset.Token.Amount.Equal(bal.Amount) {
+			println("pool Asset DS: ", asset.Token.String())
+			println("pool balance: ", bal.String())
+			return false
+		}
+	}
+
+	return true
+}
diff --git a/x/amm/keeper/msg_server_create_pool_test.go b/x/amm/keeper/msg_server_create_pool_test.go
@@ -252,6 +252,7 @@ func (suite *AmmKeeperTestSuite) TestMsgServerCreatePool() {
 			} else {
 				suite.Require().NoError(err)
 				suite.Require().Equal(resp.PoolID, uint64(1))
+				suite.Require().True(suite.VerifyPoolAssetWithBalance(1))
 
 				pools := suite.app.AmmKeeper.GetAllPool(suite.ctx)
 				suite.Require().Len(pools, 1)

diff --git a/x/amm/keeper/msg_server_exit_pool_test.go b/x/amm/keeper/msg_server_exit_pool_test.go
@@ -1,9 +1,10 @@
 package keeper_test
 
 import (
+	"time"
+
 	authtypes "github.com/cosmos/cosmos-sdk/x/auth/types"
 	govtypes "github.com/cosmos/cosmos-sdk/x/gov/types"
-	"time"
 
 	sdkmath "cosmossdk.io/math"
 	sdk "github.com/cosmos/cosmos-sdk/types"
@@ -113,14 +114,15 @@ func (suite *AmmKeeperTestSuite) TestMsgServerExitPool() {
 					Weight: sdkmath.NewInt(10),
 				},
 			}
-			_, err = msgServer.CreatePool(
+			res, err := msgServer.CreatePool(
 				suite.ctx,
 				&types.MsgCreatePool{
 					Sender:     sender.String(),
 					PoolParams: tc.poolParams,
 					PoolAssets: poolAssets,
 				})
 			suite.Require().NoError(err)
+			suite.Require().True(suite.VerifyPoolAssetWithBalance(res.PoolID))
 			pool := suite.app.AmmKeeper.GetAllPool(suite.ctx)[0]
 			suite.ctx = suite.ctx.WithBlockTime(suite.ctx.BlockTime().Add(time.Hour))
 			resp, err := msgServer.ExitPool(
@@ -137,6 +139,7 @@ func (suite *AmmKeeperTestSuite) TestMsgServerExitPool() {
 			} else {
 				suite.Require().NoError(err)
 				suite.Require().Equal(sdk.Coins(resp.TokenOut).String(), tc.minAmountsOut.String())
+				suite.Require().True(suite.VerifyPoolAssetWithBalance(pool.PoolId))
 
 				pools := suite.app.AmmKeeper.GetAllPool(suite.ctx)
 				suite.Require().Len(pools, 1)

diff --git a/x/amm/keeper/msg_server_join_pool_test.go b/x/amm/keeper/msg_server_join_pool_test.go
@@ -143,6 +143,7 @@ func (suite *AmmKeeperTestSuite) TestMsgServerJoinPool() {
 				TotalWeight: math.ZeroInt(),
 			}
 			suite.app.AmmKeeper.SetPool(suite.ctx, pool)
+			suite.Require().True(suite.VerifyPoolAssetWithBalance(1))
 
 			// execute function
 			msgServer := keeper.NewMsgServerImpl(*suite.app.AmmKeeper)
@@ -160,6 +161,7 @@ func (suite *AmmKeeperTestSuite) TestMsgServerJoinPool() {
 				suite.Require().NoError(err)
 				suite.Require().Equal(resp.ShareAmountOut, tc.shareOutAmount)
 				suite.Require().Equal(sdk.Coins(resp.TokenIn).String(), tc.expTokenIn.String())
+				suite.Require().True(suite.VerifyPoolAssetWithBalance(1))
 
 				pools := suite.app.AmmKeeper.GetAllPool(suite.ctx)
 				suite.Require().Len(pools, 1)
@@ -257,6 +259,7 @@ func (suite *AmmKeeperTestSuite) TestMsgServerJoinPoolExploitScenario() {
 				TotalWeight: math.ZeroInt(),
 			}
 			suite.app.AmmKeeper.SetPool(suite.ctx, pool)
+			suite.Require().True(suite.VerifyPoolAssetWithBalance(1))
 
 			// Step 4: Simulate market price movement - adjust weights to 10:1
 			pool.PoolAssets[0].Weight = math.NewInt(10)
@@ -275,6 +278,7 @@ func (suite *AmmKeeperTestSuite) TestMsgServerJoinPoolExploitScenario() {
 				})
 
 			suite.Require().NoError(err)
+			suite.Require().True(suite.VerifyPoolAssetWithBalance(1))
 
 			// Step 6: Validate if exploit was successful (It should fail)
 			// Calculate expected number of shares without weight balance bonus

diff --git a/x/amm/keeper/msg_server_swap_by_denom_test.go b/x/amm/keeper/msg_server_swap_by_denom_test.go
@@ -54,7 +54,7 @@ func (suite *AmmKeeperTestSuite) TestMsgServerSwapByDenom() {
 		},
 	} {
 		suite.Run(tc.desc, func() {
-			//suite.SetupTest()
+			suite.SetupTest()
 
 			// set asset profile
 			suite.app.AssetprofileKeeper.SetEntry(suite.ctx, assetprofiletypes.Entry{
@@ -150,6 +150,8 @@ func (suite *AmmKeeperTestSuite) TestMsgServerSwapByDenom() {
 			}
 			suite.app.AmmKeeper.SetPool(suite.ctx, pool)
 			suite.app.AmmKeeper.SetPool(suite.ctx, pool2)
+			suite.Require().True(suite.VerifyPoolAssetWithBalance(1))
+			suite.Require().True(suite.VerifyPoolAssetWithBalance(2))
 
 			msgServer := keeper.NewMsgServerImpl(*suite.app.AmmKeeper)
 			resp, err := msgServer.SwapByDenom(
@@ -168,6 +170,8 @@ func (suite *AmmKeeperTestSuite) TestMsgServerSwapByDenom() {
 				suite.Require().NoError(err)
 				suite.Require().Equal(resp.Amount.String(), tc.tokenOut.String())
 				suite.app.AmmKeeper.EndBlocker(suite.ctx)
+				suite.Require().True(suite.VerifyPoolAssetWithBalance(1))
+				suite.Require().True(suite.VerifyPoolAssetWithBalance(2))
 
 				// check balance change on sender
 				balances := suite.app.BankKeeper.GetAllBalances(suite.ctx, sender)

diff --git a/x/amm/keeper/msg_server_swap_exact_amount_in_test.go b/x/amm/keeper/msg_server_swap_exact_amount_in_test.go
@@ -177,6 +177,8 @@ func (suite *AmmKeeperTestSuite) TestMsgServerSwapExactAmountIn() {
 			}
 			suite.app.AmmKeeper.SetPool(suite.ctx, pool)
 			suite.app.AmmKeeper.SetPool(suite.ctx, pool2)
+			suite.Require().True(suite.VerifyPoolAssetWithBalance(1))
+			suite.Require().True(suite.VerifyPoolAssetWithBalance(2))
 
 			msgServer := keeper.NewMsgServerImpl(*suite.app.AmmKeeper)
 			resp, err := msgServer.SwapExactAmountIn(
@@ -193,6 +195,8 @@ func (suite *AmmKeeperTestSuite) TestMsgServerSwapExactAmountIn() {
 				suite.Require().NoError(err)
 				suite.Require().Equal(resp.TokenOutAmount.String(), tc.tokenOut.Amount.String())
 				suite.app.AmmKeeper.EndBlocker(suite.ctx)
+				suite.Require().True(suite.VerifyPoolAssetWithBalance(1))
+				suite.Require().True(suite.VerifyPoolAssetWithBalance(2))
 
 				// check balance change on sender
 				balances := suite.app.BankKeeper.GetAllBalances(suite.ctx, sender)
@@ -270,6 +274,7 @@ func (suite *AmmKeeperTestSuite) TestMsgServerSlippageDifferenceWhenSplit() {
 		TotalWeight: sdkmath.ZeroInt(),
 	}
 	suite.app.AmmKeeper.SetPool(suite.ctx, pool)
+	suite.Require().True(suite.VerifyPoolAssetWithBalance(1))
 
 	cacheCtx, _ := suite.ctx.CacheContext()
 	msgServer := keeper.NewMsgServerImpl(*suite.app.AmmKeeper)
@@ -284,6 +289,7 @@ func (suite *AmmKeeperTestSuite) TestMsgServerSlippageDifferenceWhenSplit() {
 	suite.Require().NoError(err)
 	suite.Require().Equal(tokenOut.Amount.String(), resp.TokenOutAmount.String())
 	suite.app.AmmKeeper.EndBlocker(cacheCtx)
+	suite.Require().True(suite.VerifyPoolAssetWithBalance(1))
 
 	// check balance change on sender
 	balances := suite.app.BankKeeper.GetAllBalances(cacheCtx, sender)
@@ -301,10 +307,12 @@ func (suite *AmmKeeperTestSuite) TestMsgServerSlippageDifferenceWhenSplit() {
 				TokenOutMinAmount: tokenOutMin,
 			})
 		suite.Require().NoError(err)
+		suite.Require().True(suite.VerifyPoolAssetWithBalance(1))
 		fmt.Printf("outAmount%d: %s\n", i, resp.TokenOutAmount.String())
 	}
 	suite.app.AmmKeeper.EndBlocker(cacheCtx)
 	// check balance change on sender after splitting swap to 100
 	balances = suite.app.BankKeeper.GetAllBalances(cacheCtx, sender)
 	suite.Require().Equal(balances.String(), expSenderBalanceSplitSwap.String())
+	suite.Require().True(suite.VerifyPoolAssetWithBalance(1))
 }
diff --git a/x/amm/keeper/msg_server_swap_exact_amount_out_test.go b/x/amm/keeper/msg_server_swap_exact_amount_out_test.go
@@ -174,6 +174,8 @@ func (suite *AmmKeeperTestSuite) TestMsgServerSwapExactAmountOut() {
 			}
 			suite.app.AmmKeeper.SetPool(suite.ctx, pool)
 			suite.app.AmmKeeper.SetPool(suite.ctx, pool2)
+			suite.Require().True(suite.VerifyPoolAssetWithBalance(1))
+			suite.Require().True(suite.VerifyPoolAssetWithBalance(2))
 
 			msgServer := keeper.NewMsgServerImpl(*suite.app.AmmKeeper)
 			resp, err := msgServer.SwapExactAmountOut(
@@ -190,6 +192,8 @@ func (suite *AmmKeeperTestSuite) TestMsgServerSwapExactAmountOut() {
 				suite.Require().NoError(err)
 				suite.Require().Equal(resp.TokenInAmount.String(), tc.tokenIn.Amount.String())
 				suite.app.AmmKeeper.EndBlocker(suite.ctx)
+				suite.Require().True(suite.VerifyPoolAssetWithBalance(1))
+				suite.Require().True(suite.VerifyPoolAssetWithBalance(2))
 
 				// check balance change on sender
 				balances := suite.app.BankKeeper.GetAllBalances(suite.ctx, sender)

diff --git a/x/amm/keeper/route_exact_amount_in_test.go b/x/amm/keeper/route_exact_amount_in_test.go
@@ -174,6 +174,7 @@ func (suite *AmmKeeperTestSuite) TestRouteExactAmountIn() {
 				TotalWeight: sdkmath.ZeroInt(),
 			}
 			suite.app.AmmKeeper.SetPool(suite.ctx, pool)
+			suite.Require().True(suite.VerifyPoolAssetWithBalance(1))
 
 			// TODO: add multiple route case
 			// TODO: add invalid route case
@@ -195,6 +196,7 @@ func (suite *AmmKeeperTestSuite) TestRouteExactAmountIn() {
 			} else {
 				suite.Require().NoError(err)
 				suite.Require().Equal(tokenOut.String(), tc.tokenOut.Amount.String())
+				suite.Require().True(suite.VerifyPoolAssetWithBalance(1))
 
 				// check pool balance increase/decrease
 				balances := suite.app.BankKeeper.GetAllBalances(suite.ctx, poolAddr)

diff --git a/x/amm/keeper/route_exact_amount_out_test.go b/x/amm/keeper/route_exact_amount_out_test.go
@@ -173,6 +173,7 @@ func (suite *AmmKeeperTestSuite) TestRouteExactAmountOut() {
 				TotalWeight: math.ZeroInt(),
 			}
 			suite.app.AmmKeeper.SetPool(suite.ctx, pool)
+			suite.Require().True(suite.VerifyPoolAssetWithBalance(1))
 
 			// TODO: add multiple route case
 			// TODO: add invalid route case
@@ -190,6 +191,7 @@ func (suite *AmmKeeperTestSuite) TestRouteExactAmountOut() {
 			} else {
 				suite.Require().NoError(err)
 				suite.Require().Equal(tokenInAmount.String(), tc.tokenIn.Amount.String())
+				suite.Require().True(suite.VerifyPoolAssetWithBalance(1))
 
 				// check pool balance increase/decrease
 				balances := suite.app.BankKeeper.GetAllBalances(suite.ctx, poolAddr)

diff --git a/x/amm/keeper/update_pool_for_swap_test.go b/x/amm/keeper/update_pool_for_swap_test.go
@@ -160,11 +160,15 @@ func (suite *AmmKeeperTestSuite) TestUpdatePoolForSwap() {
 				},
 				TotalWeight: sdkmath.ZeroInt(),
 			}
+			suite.app.AmmKeeper.SetPool(suite.ctx, pool)
+			suite.Require().True(suite.VerifyPoolAssetWithBalance(1))
+
 			err = suite.app.AmmKeeper.UpdatePoolForSwap(suite.ctx, pool, sender, sender, tc.tokenIn, tc.tokenOut, tc.swapFee, tc.tokenIn.Amount, tc.tokenOut.Amount, tc.weightBalanceBonus, false)
 			if !tc.expPass {
 				suite.Require().Error(err)
 			} else {
 				suite.Require().NoError(err)
+				suite.Require().True(suite.VerifyPoolAssetWithBalance(1))
 
 				// check pool balance increase/decrease
 				balances := suite.app.BankKeeper.GetAllBalances(suite.ctx, poolAddr)