Optimize BalanceQuery

CosmWasm · Sep 13, 2021 · 0589f38 · 0589f38
1 parent 3c0bca7
commit 0589f38
Show file tree

Hide file tree

Showing 3 changed files with 49 additions and 4 deletions.
diff --git a/x/wasm/keeper/query_plugins.go b/x/wasm/keeper/query_plugins.go
@@ -168,12 +168,11 @@ func BankQuerier(bankKeeper types.BankViewKeeper) func(ctx sdk.Context, request
 			if err != nil {
 				return nil, sdkerrors.Wrap(sdkerrors.ErrInvalidAddress, request.Balance.Address)
 			}
-			coins := bankKeeper.GetAllBalances(ctx, addr)
-			amount := coins.AmountOf(request.Balance.Denom)
+			coin := bankKeeper.GetBalance(ctx, addr, request.Balance.Denom)
 			res := wasmvmtypes.BalanceResponse{
 				Amount: wasmvmtypes.Coin{
-					Denom:  request.Balance.Denom,
-					Amount: amount.String(),
+					Denom:  coin.Denom,
+					Amount: coin.Amount.String(),
 				},
 			}
 			return json.Marshal(res)

diff --git a/x/wasm/keeper/query_plugins_test.go b/x/wasm/keeper/query_plugins_test.go
@@ -1,6 +1,7 @@
 package keeper
 
 import (
+	"encoding/json"
 	"github.com/CosmWasm/wasmd/x/wasm/keeper/wasmtesting"
 	"github.com/CosmWasm/wasmd/x/wasm/types"
 	wasmvmtypes "github.com/CosmWasm/wasmvm/types"
@@ -311,6 +312,31 @@ func TestIBCQuerier(t *testing.T) {
 
 }
 
+func TestBankQuerierBalance(t *testing.T) {
+	mock := bankKeeperMock{GetBalanceFn: func(ctx sdk.Context, addr sdk.AccAddress, denom string) sdk.Coin {
+		return sdk.NewCoin(denom, sdk.NewInt(1))
+	}}
+
+	ctx := sdk.Context{}
+	q := BankQuerier(mock)
+	gotBz, gotErr := q(ctx, &wasmvmtypes.BankQuery{
+		Balance: &wasmvmtypes.BalanceQuery{
+			Address: RandomBech32AccountAddress(t),
+			Denom:   "ALX",
+		},
+	})
+	require.NoError(t, gotErr)
+	var got wasmvmtypes.BalanceResponse
+	require.NoError(t, json.Unmarshal(gotBz, &got))
+	exp := wasmvmtypes.BalanceResponse{
+		Amount: wasmvmtypes.Coin{
+			Denom:  "ALX",
+			Amount: "1",
+		},
+	}
+	assert.Equal(t, exp, got)
+}
+
 type wasmKeeperMock struct {
 	GetContractInfoFn func(ctx sdk.Context, contractAddress sdk.AccAddress) *types.ContractInfo
 }
@@ -325,3 +351,22 @@ func (m wasmKeeperMock) GetContractInfo(ctx sdk.Context, contractAddress sdk.Acc
 	}
 	return m.GetContractInfoFn(ctx, contractAddress)
 }
+
+type bankKeeperMock struct {
+	GetBalanceFn     func(ctx sdk.Context, addr sdk.AccAddress, denom string) sdk.Coin
+	GetAllBalancesFn func(ctx sdk.Context, addr sdk.AccAddress) sdk.Coins
+}
+
+func (m bankKeeperMock) GetBalance(ctx sdk.Context, addr sdk.AccAddress, denom string) sdk.Coin {
+	if m.GetBalanceFn == nil {
+		panic("not expected to be called")
+	}
+	return m.GetBalanceFn(ctx, addr, denom)
+}
+
+func (m bankKeeperMock) GetAllBalances(ctx sdk.Context, addr sdk.AccAddress) sdk.Coins {
+	if m.GetAllBalancesFn == nil {
+		panic("not expected to be called")
+	}
+	return m.GetAllBalancesFn(ctx, addr)
+}
diff --git a/x/wasm/types/expected_keepers.go b/x/wasm/types/expected_keepers.go
@@ -15,6 +15,7 @@ import (
 // BankViewKeeper defines a subset of methods implemented by the cosmos-sdk bank keeper
 type BankViewKeeper interface {
 	GetAllBalances(ctx sdk.Context, addr sdk.AccAddress) sdk.Coins
+	GetBalance(ctx sdk.Context, addr sdk.AccAddress, denom string) sdk.Coin
 }
 
 // Burner is a subset of the sdk bank keeper methods