test

cosmos · Sep 26, 2024 · f19996e · f19996e
1 parent 42f4dfd
commit f19996e
Showing 1 changed file with 52 additions and 0 deletions.
diff --git a/x/bank/v2/keeper/keeper_test.go b/x/bank/v2/keeper/keeper_test.go
@@ -1,7 +1,9 @@
 package keeper_test
 
 import (
+	"bytes"
 	"context"
+	"fmt"
 	"testing"
 	"time"
 
@@ -184,3 +186,53 @@ func (suite *KeeperTestSuite) TestSendCoins_Module_To_Module() {
 	mintBarBalance := suite.bankKeeper.GetBalance(ctx, mintAcc.GetAddress(), barDenom)
 	require.Equal(mintBarBalance.Amount, math.NewInt(0))
 }
+
+func (suite *KeeperTestSuite) TestSendCoins_WithRestriction() {
+	ctx := suite.ctx
+	require := suite.Require()
+	balances := sdk.NewCoins(newFooCoin(100), newBarCoin(50))
+	sendAmt := sdk.NewCoins(newFooCoin(10), newBarCoin(10))
+
+	require.NoError(banktestutil.FundAccount(ctx, suite.bankKeeper, accAddrs[0], balances))
+
+	// Add first restriction
+	addrRestrictFunc := func(ctx context.Context, from, to []byte, amount sdk.Coins) ([]byte, error) {
+		if bytes.Equal(from, to) {
+			return nil, fmt.Errorf("Can not send to same address")
+		}
+		return to, nil
+	}
+	suite.bankKeeper.AppendSendRestriction(addrRestrictFunc)
+
+	err := suite.bankKeeper.SendCoins(ctx, accAddrs[0], accAddrs[0], sendAmt)
+	require.Error(err)
+	require.Contains(err.Error(), "Can not send to same address")
+
+	// Add second restriction
+	amtRestrictFunc := func(ctx context.Context, from, to []byte, amount sdk.Coins) ([]byte, error) {
+		if len(amount) > 1 {
+			return nil, fmt.Errorf("Allow only one denom per one send")
+		}
+		return to, nil
+	}
+	suite.bankKeeper.AppendSendRestriction(amtRestrictFunc)
+
+	// Pass the 1st but failt at the 2nd
+	err = suite.bankKeeper.SendCoins(ctx, accAddrs[0], accAddrs[1], sendAmt)
+	require.Error(err)
+	require.Contains(err.Error(), "Allow only one denom per one send")
+
+	// Pass both 2 restrictions
+	err = suite.bankKeeper.SendCoins(ctx, accAddrs[0], accAddrs[1], sdk.NewCoins(newFooCoin(10)))
+	require.NoError(err)
+
+	// Check balances
+	acc0FooBalance := suite.bankKeeper.GetBalance(ctx, accAddrs[0], fooDenom)
+	require.Equal(acc0FooBalance.Amount, math.NewInt(90))
+	acc0BarBalance := suite.bankKeeper.GetBalance(ctx, accAddrs[0], barDenom)
+	require.Equal(acc0BarBalance.Amount, math.NewInt(50))
+	acc1FooBalance := suite.bankKeeper.GetBalance(ctx, accAddrs[1], fooDenom)
+	require.Equal(acc1FooBalance.Amount, math.NewInt(10))
+	acc1BarBalance := suite.bankKeeper.GetBalance(ctx, accAddrs[1], barDenom)
+	require.Equal(acc1BarBalance.Amount, math.ZeroInt())
+}