chore: fix unit tests of margin invariant check

elys-network · Sep 5, 2023 · 0b65a4e · 0b65a4e
1 parent 2439e94
commit 0b65a4e
Show file tree

Hide file tree

Showing 4 changed files with 22 additions and 6 deletions.
diff --git a/x/accountedpool/keeper/accounted_pool_update_test.go b/x/accountedpool/keeper/accounted_pool_update_test.go
@@ -11,6 +11,7 @@ import (
 	"github.com/elys-network/elys/x/accountedpool/types"
 	margintypes "github.com/elys-network/elys/x/margin/types"
 	ptypes "github.com/elys-network/elys/x/parameter/types"
+	"github.com/stretchr/testify/require"
 )
 
 func TestAccountedPoolUpdate(t *testing.T) {
@@ -65,8 +66,8 @@ func TestAccountedPoolUpdate(t *testing.T) {
 				AssetDenom:           ptypes.USDC,
 			},
 			{
-				Liabilities:          sdk.NewInt(50),
-				Custody:              sdk.NewInt(0),
+				Liabilities:          sdk.NewInt(0),
+				Custody:              sdk.NewInt(50),
 				AssetBalance:         sdk.NewInt(0),
 				UnsettledLiabilities: sdk.NewInt(0),
 				BlockInterest:        sdk.NewInt(0),
@@ -77,4 +78,12 @@ func TestAccountedPoolUpdate(t *testing.T) {
 	// Update accounted pool
 	apk.UpdateAccountedPool(ctx, ammPool, marginPool)
 
+	apool, found := apk.GetAccountedPool(ctx, (uint64)(0))
+	require.Equal(t, found, true)
+	require.Equal(t, apool.PoolId, (uint64)(0))
+
+	usdcBalance := apk.GetAccountedBalance(ctx, (uint64)(0), ptypes.USDC)
+	require.Equal(t, usdcBalance, sdk.NewInt(1000+400+100))
+	atomBalance := apk.GetAccountedBalance(ctx, (uint64)(0), ptypes.ATOM)
+	require.Equal(t, atomBalance, sdk.NewInt(100))
 }
diff --git a/x/margin/keeper/invariant_check.go b/x/margin/keeper/invariant_check.go
@@ -43,7 +43,7 @@ func (k Keeper) InvariantCheck(ctx sdk.Context) error {
 		ammPoolId := mtp.AmmPoolId
 		err := k.AmmPoolBalanceCheck(ctx, ammPoolId)
 		if err != nil {
-			panic(err)
+			return err
 		}
 	}
 

diff --git a/x/margin/keeper/invariant_check_test.go b/x/margin/keeper/invariant_check_test.go
@@ -1,6 +1,7 @@
 package keeper_test
 
 import (
+	"errors"
 	"testing"
 
 	tmproto "github.com/cometbft/cometbft/proto/tendermint/types"
@@ -113,7 +114,7 @@ func TestCheckBalanceInvariant_InvalidBalance(t *testing.T) {
 
 	// Check balance invariant check
 	err = mk.InvariantCheck(ctx)
-	require.NoError(t, err)
+	require.Equal(t, err, errors.New("balance mismatch!"))
 
 	mtpId := mtps[0].Id
 	// Create a margin position close msg
@@ -126,7 +127,7 @@ func TestCheckBalanceInvariant_InvalidBalance(t *testing.T) {
 	require.NoError(t, err)
 
 	balances = app.BankKeeper.GetAllBalances(ctx, poolAddress)
-	require.Equal(t, balances.AmountOf(ptypes.USDC), sdk.NewInt(10000))
+	require.Equal(t, balances.AmountOf(ptypes.USDC), sdk.NewInt(10046))
 	require.Equal(t, balances.AmountOf(ptypes.ATOM), sdk.NewInt(100000))
 
 	// Check balance invariant check

diff --git a/x/margin/keeper/keeper.go b/x/margin/keeper/keeper.go
@@ -492,7 +492,13 @@ func (k Keeper) Repay(ctx sdk.Context, mtp *types.MTP, pool *types.Pool, ammPool
 			if err != nil {
 				return err
 			}
-			err = k.bankKeeper.SendCoinsFromModuleToAccount(ctx, ammPool.Address, addr, returnCoins)
+
+			ammPoolAddr, err := sdk.AccAddressFromBech32(ammPool.Address)
+			if err != nil {
+				return err
+			}
+
+			err = k.bankKeeper.SendCoins(ctx, ammPoolAddr, addr, returnCoins)
 			if err != nil {
 				return err
 			}