refactor: margin functions (#194)

* refactor: margin functions

* fix: test cases

* fix: mocks

* test: fix test cases

x/margin/keeper/check_long_assets.go

@@ -7,7 +7,7 @@ import (
 	ptypes "github.com/elys-network/elys/x/parameter/types"
 )
 
-func (k Keeper) CheckLongingAssets(ctx sdk.Context, collateralAsset string, borrowAsset string) error {
+func (k Keeper) CheckLongAssets(ctx sdk.Context, collateralAsset string, borrowAsset string) error {
 	if borrowAsset == ptypes.BaseCurrency {
 		return sdkerrors.Wrap(types.ErrInvalidBorrowingAsset, "invalid borrowing asset")
 	}

x/margin/keeper/check_long_assets_test.go

@@ -25,10 +25,10 @@ func TestCheckLongAssets_InvalidAssets(t *testing.T) {
 
 	ctx := sdk.Context{} // mock or setup a context
 
-	err := k.CheckLongingAssets(ctx, ptypes.BaseCurrency, ptypes.BaseCurrency)
+	err := k.CheckLongAssets(ctx, ptypes.BaseCurrency, ptypes.BaseCurrency)
 	assert.True(t, errors.Is(err, sdkerrors.Wrap(types.ErrInvalidBorrowingAsset, "invalid borrowing asset")))
 
-	err = k.CheckLongingAssets(ctx, ptypes.ATOM, ptypes.BaseCurrency)
+	err = k.CheckLongAssets(ctx, ptypes.ATOM, ptypes.BaseCurrency)
 	assert.True(t, errors.Is(err, sdkerrors.Wrap(types.ErrInvalidBorrowingAsset, "invalid borrowing asset")))
 
 	// Expect no error
@@ -46,10 +46,10 @@ func TestCheckLongAssets_ValidAssets(t *testing.T) {
 
 	ctx := sdk.Context{} // mock or setup a context
 
-	err := k.CheckLongingAssets(ctx, ptypes.BaseCurrency, ptypes.ATOM)
+	err := k.CheckLongAssets(ctx, ptypes.BaseCurrency, ptypes.ATOM)
 	assert.Nil(t, err)
 
-	err = k.CheckLongingAssets(ctx, ptypes.ATOM, ptypes.ATOM)
+	err = k.CheckLongAssets(ctx, ptypes.ATOM, ptypes.ATOM)
 	assert.Nil(t, err)
 
 	// Expect an error about max open positions

x/margin/keeper/check_short_assets.go

@@ -0,0 +1,27 @@
+package keeper
+
+import (
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	sdkerrors "github.com/cosmos/cosmos-sdk/types/errors"
+	"github.com/elys-network/elys/x/margin/types"
+	ptypes "github.com/elys-network/elys/x/parameter/types"
+)
+
+func (k Keeper) CheckShortAssets(ctx sdk.Context, collateralAsset string, borrowAsset string) error {
+	// You shouldn't be shorting the base currency (like USDC).
+	if borrowAsset == ptypes.BaseCurrency {
+		return sdkerrors.Wrap(types.ErrInvalidBorrowingAsset, "cannot short the base currency")
+	}
+
+	// If both the collateralAsset and borrowAsset are the same, it doesn't make sense.
+	if collateralAsset == borrowAsset {
+		return sdkerrors.Wrap(types.ErrInvalidCollateralAsset, "collateral asset cannot be the same as the borrowed asset in a short position")
+	}
+
+	// The collateral for a short must be the base currency.
+	if collateralAsset != ptypes.BaseCurrency {
+		return sdkerrors.Wrap(types.ErrInvalidCollateralAsset, "collateral asset for a short position must be the base currency")
+	}
+
+	return nil
+}

x/margin/keeper/check_short_assets_test.go

@@ -0,0 +1,59 @@
+package keeper_test
+
+import (
+	"errors"
+	"testing"
+
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	sdkerrors "github.com/cosmos/cosmos-sdk/types/errors"
+	"github.com/elys-network/elys/x/margin/keeper"
+	"github.com/elys-network/elys/x/margin/types"
+	"github.com/elys-network/elys/x/margin/types/mocks"
+	"github.com/stretchr/testify/assert"
+
+	ptypes "github.com/elys-network/elys/x/parameter/types"
+)
+
+func TestCheckShortAssets_InvalidAssets(t *testing.T) {
+	// Setup the mock checker
+	mockChecker := new(mocks.OpenShortChecker)
+
+	// Create an instance of Keeper with the mock checker
+	k := keeper.Keeper{
+		OpenShortChecker: mockChecker,
+	}
+
+	ctx := sdk.Context{} // mock or setup a context
+
+	// Test invalid cases for short positions
+	err := k.CheckShortAssets(ctx, ptypes.ATOM, ptypes.BaseCurrency)
+	assert.True(t, errors.Is(err, sdkerrors.Wrap(types.ErrInvalidBorrowingAsset, "cannot short the base currency")))
+
+	err = k.CheckShortAssets(ctx, ptypes.ATOM, ptypes.ATOM)
+	assert.True(t, errors.Is(err, sdkerrors.Wrap(types.ErrInvalidCollateralAsset, "collateral asset cannot be the same as the borrowed asset in a short position")))
+
+	err = k.CheckShortAssets(ctx, ptypes.ATOM, "btc")
+	assert.True(t, errors.Is(err, sdkerrors.Wrap(types.ErrInvalidCollateralAsset, "collateral asset for a short position must be the base currency")))
+
+	// Expect no error
+	mockChecker.AssertExpectations(t)
+}
+
+func TestCheckShortAssets_ValidAssets(t *testing.T) {
+	// Setup the mock checker
+	mockChecker := new(mocks.OpenShortChecker)
+
+	// Create an instance of Keeper with the mock checker
+	k := keeper.Keeper{
+		OpenShortChecker: mockChecker,
+	}
+
+	ctx := sdk.Context{} // mock or setup a context
+
+	// Test valid case for short position
+	err := k.CheckShortAssets(ctx, ptypes.BaseCurrency, ptypes.ATOM)
+	assert.Nil(t, err)
+
+	// Expect no error
+	mockChecker.AssertExpectations(t)
+}

x/margin/keeper/open.go

@@ -7,8 +7,18 @@ import (
 )
 
 func (k Keeper) Open(ctx sdk.Context, msg *types.MsgOpen) (*types.MsgOpenResponse, error) {
-	if err := k.OpenChecker.CheckLongingAssets(ctx, msg.CollateralAsset, msg.BorrowAsset); err != nil {
-		return nil, err
+	// Determine the type of position (long or short) and validate assets accordingly.
+	switch msg.Position {
+	case types.Position_LONG:
+		if err := k.OpenChecker.CheckLongAssets(ctx, msg.CollateralAsset, msg.BorrowAsset); err != nil {
+			return nil, err
+		}
+	case types.Position_SHORT:
+		if err := k.OpenChecker.CheckShortAssets(ctx, msg.CollateralAsset, msg.BorrowAsset); err != nil {
+			return nil, err
+		}
+	default:
+		return nil, sdkerrors.Wrap(types.ErrInvalidPosition, msg.Position.String())
 	}
 
 	if err := k.OpenChecker.CheckUserAuthorization(ctx, msg); err != nil {

x/margin/keeper/open_long_process.go

@@ -8,24 +8,29 @@ import (
 )
 
 func (k Keeper) ProcessOpenLong(ctx sdk.Context, mtp *types.MTP, leverage sdk.Dec, eta sdk.Dec, collateralAmountDec sdk.Dec, poolId uint64, msg *types.MsgOpen) (*types.MTP, error) {
-	// Get token asset other than base currency
+	// Determine the trading asset.
 	tradingAsset := k.OpenLongChecker.GetTradingAsset(msg.CollateralAsset, msg.BorrowAsset)
 
+	// Fetch the pool associated with the given pool ID.
 	pool, found := k.OpenLongChecker.GetPool(ctx, poolId)
 	if !found {
 		return nil, sdkerrors.Wrap(types.ErrPoolDoesNotExist, tradingAsset)
 	}
 
+	// Check if the pool is enabled.
 	if !k.OpenLongChecker.IsPoolEnabled(ctx, poolId) {
 		return nil, sdkerrors.Wrap(types.ErrMTPDisabled, tradingAsset)
 	}
 
+	// Fetch the corresponding AMM (Automated Market Maker) pool.
 	ammPool, err := k.OpenLongChecker.GetAmmPool(ctx, poolId, tradingAsset)
 	if err != nil {
 		return nil, err
 	}
 
+	// Calculate the leveraged amount based on the collateral provided and the leverage.
 	leveragedAmount := sdk.NewInt(collateralAmountDec.Mul(leverage).TruncateInt().Int64())
+
 	// If collateral is not base currency, calculate the borrowing amount in base currency and check the balance
 	if msg.CollateralAsset != ptypes.BaseCurrency {
 		custodyAmtToken := sdk.NewCoin(msg.CollateralAsset, leveragedAmount)
@@ -42,12 +47,14 @@ func (k Keeper) ProcessOpenLong(ctx sdk.Context, mtp *types.MTP, leverage sdk.De
 		}
 	}
 
+	// Check minimum liabilities.
 	collateralTokenAmt := sdk.NewCoin(msg.CollateralAsset, msg.CollateralAmount)
 	err = k.OpenLongChecker.CheckMinLiabilities(ctx, collateralTokenAmt, eta, pool, ammPool, msg.BorrowAsset)
 	if err != nil {
 		return nil, err
 	}
 
+	// Calculate custody amount.
 	leveragedAmtTokenIn := sdk.NewCoin(msg.CollateralAsset, leveragedAmount)
 	custodyAmount, err := k.OpenLongChecker.EstimateSwap(ctx, leveragedAmtTokenIn, msg.BorrowAsset, ammPool)
 	if err != nil {
@@ -59,26 +66,34 @@ func (k Keeper) ProcessOpenLong(ctx sdk.Context, mtp *types.MTP, leverage sdk.De
 		custodyAmount = leveragedAmount
 	}
 
+	// Ensure the AMM pool has enough balance.
 	if !k.OpenLongChecker.HasSufficientPoolBalance(ctx, ammPool, msg.BorrowAsset, custodyAmount) {
 		return nil, sdkerrors.Wrap(types.ErrCustodyTooHigh, custodyAmount.String())
 	}
 
+	// Borrow the asset the user wants to long.
 	err = k.OpenLongChecker.Borrow(ctx, msg.CollateralAsset, msg.BorrowAsset, msg.CollateralAmount, custodyAmount, mtp, &ammPool, &pool, eta)
 	if err != nil {
 		return nil, err
 	}
+
+	// Update the pool health.
 	if err = k.OpenLongChecker.UpdatePoolHealth(ctx, &pool); err != nil {
 		return nil, err
 	}
+
+	// Take custody from the pool balance.
 	if err = k.OpenLongChecker.TakeInCustody(ctx, *mtp, &pool); err != nil {
 		return nil, err
 	}
 
+	// Update the MTP health.
 	lr, err := k.OpenLongChecker.UpdateMTPHealth(ctx, *mtp, ammPool)
 	if err != nil {
 		return nil, err
 	}
 
+	// Check if the MTP is unhealthy
 	safetyFactor := k.OpenLongChecker.GetSafetyFactor(ctx)
 	if lr.LTE(safetyFactor) {
 		return nil, types.ErrMTPUnhealthy

x/margin/keeper/open_short_process.go

@@ -6,7 +6,7 @@ import (
 )
 
 func (k Keeper) ProcessOpenShort(ctx sdk.Context, mtp *types.MTP, leverage sdk.Dec, eta sdk.Dec, collateralAmountDec sdk.Dec, poolId uint64, msg *types.MsgOpen) (*types.MTP, error) {
-	// // Determine the trading asset.
+	// Determine the trading asset.
 	// tradingAsset := k.OpenShortChecker.GetTradingAsset(msg.CollateralAsset, msg.BorrowAsset)
 
 	// // Fetch the pool associated with the given pool ID.
@@ -30,9 +30,12 @@ func (k Keeper) ProcessOpenShort(ctx sdk.Context, mtp *types.MTP, leverage sdk.D
 	// leveragedAmount := sdk.NewInt(collateralAmountDec.Mul(leverage).TruncateInt().Int64())
 
 	// // Borrow the asset the user wants to short.
-	// // ... (Logic to borrow the asset; error handling) ...
+	// // err = k.OpenLongChecker.Borrow(ctx, msg.CollateralAsset, msg.BorrowAsset, msg.CollateralAmount, custodyAmount, mtp, &ammPool, &pool, eta)
+	// // if err != nil {
+	// // 	return nil, err
+	// // }
 
-	// // Swap the borrowed asset for base currency.
+	// // Calculate the custody amount.
 	// swappedAmount, err := k.OpenShortChecker.EstimateSwap(ctx, leveragedAmount, ptypes.BaseCurrency, ammPool)
 	// if err != nil {
 	// 	return nil, err