From f19996ece56e7d50b6086c209ac7d567f6d00afc Mon Sep 17 00:00:00 2001
From: Hieu Vu <72878483+hieuvubk@users.noreply.github.com>
Date: Thu, 26 Sep 2024 15:13:48 +0700
Subject: [PATCH] test

---
 x/bank/v2/keeper/keeper_test.go | 52 +++++++++++++++++++++++++++++++++
 1 file changed, 52 insertions(+)

diff --git a/x/bank/v2/keeper/keeper_test.go b/x/bank/v2/keeper/keeper_test.go
index 32ccbac4d2ef..4da9faed3b5c 100644
--- 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())
+}