feat(x/precisebank): Return full balances in GetBalance(), add SpendableCoin method (#1957)

Change GetBalance() to return full balances instead of spendable to align behavior with x/bank. Add SpendableCoin() method with support of akava for use in x/evm.

x/precisebank/keeper/view.go

@@ -18,15 +18,42 @@ func (k Keeper) GetBalance(
 		return k.bk.GetBalance(ctx, addr, denom)
 	}
 
-	// x/bank for integer balance - spendable balance only
+	// x/bank for integer balance - full balance, including locked
-	spendableCoins := k.bk.SpendableCoins(ctx, addr)
+	integerCoins := k.bk.GetBalance(ctx, addr, types.IntegerCoinDenom)
-	integerAmount := spendableCoins.AmountOf(types.IntegerCoinDenom)
 
 	// x/precisebank for fractional balance
 	fractionalAmount := k.GetFractionalBalance(ctx, addr)
 
 	// (Integer * ConversionFactor) + Fractional
-	fullAmount := integerAmount.
+	fullAmount := integerCoins.
+		Amount.
+		Mul(types.ConversionFactor()).
+		Add(fractionalAmount)
+
+	return sdk.NewCoin(types.ExtendedCoinDenom, fullAmount)
+}
+
+// SpendableCoins returns the total balances of spendable coins for an account
+// by address. If the account has no spendable coins, an empty Coins slice is
+// returned.
+func (k Keeper) SpendableCoin(
+	ctx sdk.Context,
+	addr sdk.AccAddress,
+	denom string,
+) sdk.Coin {
+	// Pass through to x/bank for denoms except ExtendedCoinDenom
+	if denom != types.ExtendedCoinDenom {
+		return k.bk.SpendableCoin(ctx, addr, denom)
+	}
+
+	// x/bank for integer balance - excluding locked
+	integerCoin := k.bk.SpendableCoin(ctx, addr, types.IntegerCoinDenom)
+
+	// x/precisebank for fractional balance
+	fractionalAmount := k.GetFractionalBalance(ctx, addr)
+
+	// Spendable = (Integer * ConversionFactor) + Fractional
+	fullAmount := integerCoin.Amount.
 		Mul(types.ConversionFactor()).
 		Add(fractionalAmount)
 

x/precisebank/keeper/view_integration_test.go

@@ -0,0 +1,131 @@
+package keeper_test
+
+import (
+	"testing"
+
+	sdkmath "cosmossdk.io/math"
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	authtypes "github.com/cosmos/cosmos-sdk/x/auth/types"
+	vestingtypes "github.com/cosmos/cosmos-sdk/x/auth/vesting/types"
+
+	"github.com/kava-labs/kava/x/precisebank/testutil"
+	"github.com/kava-labs/kava/x/precisebank/types"
+	"github.com/stretchr/testify/suite"
+)
+
+type viewIntegrationTestSuite struct {
+	testutil.Suite
+}
+
+func (suite *viewIntegrationTestSuite) SetupTest() {
+	suite.Suite.SetupTest()
+}
+
+func TestViewIntegrationTestSuite(t *testing.T) {
+	suite.Run(t, new(viewIntegrationTestSuite))
+}
+
+func (suite *viewIntegrationTestSuite) TestKeeper_SpendableCoin() {
+	tests := []struct {
+		name      string
+		giveDenom string // queried denom for balance
+
+		giveBankBal       sdk.Coins   // full balance
+		giveFractionalBal sdkmath.Int // stored fractional balance for giveAddr
+		giveLockedCoins   sdk.Coins   // locked coins
+
+		wantSpendableBal sdk.Coin
+	}{
+		{
+			"extended denom, no fractional - locked coins",
+			types.ExtendedCoinDenom,
+			// queried bank balance in ukava when querying for akava
+			sdk.NewCoins(sdk.NewCoin(types.IntegerCoinDenom, sdk.NewInt(1000))),
+			sdkmath.ZeroInt(),
+			sdk.NewCoins(sdk.NewCoin(types.IntegerCoinDenom, sdk.NewInt(10))),
+			// (integer + fractional) - locked
+			sdk.NewCoin(
+				types.ExtendedCoinDenom,
+				types.ConversionFactor().MulRaw(1000-10),
+			),
+		},
+		{
+			"extended denom, with fractional - locked coins",
+			types.ExtendedCoinDenom,
+			// queried bank balance in ukava when querying for akava
+			sdk.NewCoins(sdk.NewCoin(types.IntegerCoinDenom, sdk.NewInt(1000))),
+			sdkmath.NewInt(5000),
+			sdk.NewCoins(sdk.NewCoin(types.IntegerCoinDenom, sdk.NewInt(10))),
+			sdk.NewCoin(
+				types.ExtendedCoinDenom,
+				// (integer - locked) + fractional
+				types.ConversionFactor().MulRaw(1000-10).AddRaw(5000),
+			),
+		},
+		{
+			"non-extended denom - ukava returns ukava",
+			types.IntegerCoinDenom,
+			sdk.NewCoins(sdk.NewCoin(types.IntegerCoinDenom, sdk.NewInt(1000))),
+			sdkmath.ZeroInt(),
+			sdk.NewCoins(sdk.NewCoin(types.IntegerCoinDenom, sdk.NewInt(10))),
+			sdk.NewCoin(types.IntegerCoinDenom, sdk.NewInt(990)),
+		},
+		{
+			"non-extended denom, with fractional - ukava returns ukava",
+			types.IntegerCoinDenom,
+			sdk.NewCoins(sdk.NewCoin(types.IntegerCoinDenom, sdk.NewInt(1000))),
+			// does not affect balance
+			sdkmath.NewInt(100),
+			sdk.NewCoins(sdk.NewCoin(types.IntegerCoinDenom, sdk.NewInt(10))),
+			sdk.NewCoin(types.IntegerCoinDenom, sdk.NewInt(990)),
+		},
+	}
+
+	for _, tt := range tests {
+		suite.Run(tt.name, func() {
+			suite.SetupTest()
+
+			addr := sdk.AccAddress([]byte("test-address"))
+
+			suite.MintToAccount(addr, tt.giveBankBal)
+
+			// Set fractional balance in store before query
+			suite.Keeper.SetFractionalBalance(suite.Ctx, addr, tt.giveFractionalBal)
+
+			// Add some locked coins
+			acc := suite.AccountKeeper.GetAccount(suite.Ctx, addr)
+			if acc == nil {
+				acc = authtypes.NewBaseAccount(addr, nil, 0, 0)
+			}
+
+			vestingAcc := vestingtypes.NewPeriodicVestingAccount(
+				acc.(*authtypes.BaseAccount),
+				tt.giveLockedCoins,
+				suite.Ctx.BlockTime().Unix(),
+				vestingtypes.Periods{
+					vestingtypes.Period{
+						Length: 100,
+						Amount: tt.giveLockedCoins,
+					},
+				},
+			)
+			suite.AccountKeeper.SetAccount(suite.Ctx, vestingAcc)
+
+			fetchedLockedCoins := vestingAcc.LockedCoins(suite.Ctx.BlockTime())
+			suite.Require().Equal(
+				tt.giveLockedCoins,
+				fetchedLockedCoins,
+				"locked coins should be matching at current block time",
+			)
+
+			spendableCoinsWithLocked := suite.Keeper.SpendableCoin(suite.Ctx, addr, tt.giveDenom)
+
+			suite.Require().Equalf(
+				tt.wantSpendableBal,
+				spendableCoinsWithLocked,
+				"expected spendable coins of denom %s",
+				tt.giveDenom,
+			)
+		})
+	}
+}

x/precisebank/keeper/view_test.go

@@ -93,9 +93,10 @@ func TestKeeper_GetBalance(t *testing.T) {
 			if tt.giveDenom == types.ExtendedCoinDenom {
 				// No balance pass through
 				tk.bk.EXPECT().
-					SpendableCoins(tk.ctx, addr).
+					GetBalance(tk.ctx, addr, types.IntegerCoinDenom).
-					RunAndReturn(func(_ sdk.Context, _ sdk.AccAddress) sdk.Coins {
+					RunAndReturn(func(_ sdk.Context, _ sdk.AccAddress, _ string) sdk.Coin {
-						return tt.giveBankBal
+						amt := tt.giveBankBal.AmountOf(types.IntegerCoinDenom)
+						return sdk.NewCoin(types.IntegerCoinDenom, amt)
 					}).
 					Once()
 			} else {
@@ -115,3 +116,111 @@ func TestKeeper_GetBalance(t *testing.T) {
 		})
 	}
 }
+
+func TestKeeper_SpendableCoin(t *testing.T) {
+	tests := []struct {
+		name      string
+		giveDenom string // queried denom for balance
+
+		giveBankBal       sdk.Coins   // mocked bank balance for giveAddr
+		giveFractionalBal sdkmath.Int // stored fractional balance for giveAddr
+
+		wantBal sdk.Coin
+	}{
+		{
+			"extended denom - no fractional balance",
+			types.ExtendedCoinDenom,
+			// queried bank balance in ukava when querying for akava
+			sdk.NewCoins(sdk.NewCoin(types.IntegerCoinDenom, sdk.NewInt(1000))),
+			sdkmath.ZeroInt(),
+			// integer + fractional
+			sdk.NewCoin(types.ExtendedCoinDenom, sdk.NewInt(1000_000_000_000_000)),
+		},
+		{
+			"extended denom - with fractional balance",
+			types.ExtendedCoinDenom,
+			sdk.NewCoins(sdk.NewCoin(types.IntegerCoinDenom, sdk.NewInt(1000))),
+			sdkmath.NewInt(100),
+			// integer + fractional
+			sdk.NewCoin(types.ExtendedCoinDenom, sdk.NewInt(1000_000_000_000_100)),
+		},
+		{
+			"extended denom - only fractional balance",
+			types.ExtendedCoinDenom,
+			// no coins in bank, only fractional balance
+			sdk.NewCoins(),
+			sdkmath.NewInt(100),
+			sdk.NewCoin(types.ExtendedCoinDenom, sdk.NewInt(100)),
+		},
+		{
+			"extended denom - max fractional balance",
+			types.ExtendedCoinDenom,
+			sdk.NewCoins(sdk.NewCoin(types.IntegerCoinDenom, sdk.NewInt(1000))),
+			types.ConversionFactor().SubRaw(1),
+			// integer + fractional
+			sdk.NewCoin(types.ExtendedCoinDenom, sdk.NewInt(1000_999_999_999_999)),
+		},
+		{
+			"non-extended denom - ukava returns ukava",
+			types.IntegerCoinDenom,
+			sdk.NewCoins(sdk.NewCoin(types.IntegerCoinDenom, sdk.NewInt(1000))),
+			sdk.ZeroInt(),
+			sdk.NewCoin("ukava", sdk.NewInt(1000)),
+		},
+		{
+			"non-extended denom - unaffected by fractional balance",
+			"ukava",
+			sdk.NewCoins(sdk.NewCoin("ukava", sdk.NewInt(1000))),
+			sdkmath.NewInt(100),
+			sdk.NewCoin("ukava", sdk.NewInt(1000)),
+		},
+		{
+			"unrelated denom - no fractional",
+			"busd",
+			sdk.NewCoins(sdk.NewCoin("busd", sdk.NewInt(1000))),
+			sdk.ZeroInt(),
+			sdk.NewCoin("busd", sdk.NewInt(1000)),
+		},
+		{
+			"unrelated denom - unaffected by fractional balance",
+			"busd",
+			sdk.NewCoins(sdk.NewCoin("busd", sdk.NewInt(1000))),
+			sdkmath.NewInt(100),
+			sdk.NewCoin("busd", sdk.NewInt(1000)),
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			tk := NewMockedTestData(t)
+			addr := sdk.AccAddress([]byte("test-address"))
+
+			// Set fractional balance in store before query
+			tk.keeper.SetFractionalBalance(tk.ctx, addr, tt.giveFractionalBal)
+
+			if tt.giveDenom == types.ExtendedCoinDenom {
+				// No balance pass through
+				tk.bk.EXPECT().
+					SpendableCoin(tk.ctx, addr, types.IntegerCoinDenom).
+					RunAndReturn(func(_ sdk.Context, _ sdk.AccAddress, _ string) sdk.Coin {
+						amt := tt.giveBankBal.AmountOf(types.IntegerCoinDenom)
+						return sdk.NewCoin(types.IntegerCoinDenom, amt)
+					}).
+					Once()
+			} else {
+				// Pass through to x/bank for denoms except ExtendedCoinDenom
+				tk.bk.EXPECT().
+					SpendableCoin(tk.ctx, addr, tt.giveDenom).
+					RunAndReturn(func(ctx sdk.Context, aa sdk.AccAddress, s string) sdk.Coin {
+						require.Equal(t, s, tt.giveDenom, "unexpected denom passed to x/bank.GetBalance")
+
+						return sdk.NewCoin(tt.giveDenom, tt.giveBankBal.AmountOf(s))
+					}).
+					Once()
+			}
+
+			bal := tk.keeper.SpendableCoin(tk.ctx, addr, tt.giveDenom)
+			require.Equal(t, tt.wantBal, bal)
+		})
+	}
+}

x/precisebank/types/expected_keepers.go

@@ -17,7 +17,7 @@ type BankKeeper interface {
 	IsSendEnabledCoins(ctx sdk.Context, coins ...sdk.Coin) error
 	GetBalance(ctx sdk.Context, addr sdk.AccAddress, denom string) sdk.Coin
 	GetSupply(ctx sdk.Context, denom string) sdk.Coin
-	SpendableCoins(ctx sdk.Context, addr sdk.AccAddress) sdk.Coins
+	SpendableCoin(ctx sdk.Context, addr sdk.AccAddress, denom string) sdk.Coin
 
 	BlockedAddr(addr sdk.AccAddress) bool
 

x/precisebank/types/mocks/MockBankKeeper.go

@@ -515,51 +515,50 @@ func (_c *MockBankKeeper_SendCoinsFromModuleToModule_Call) RunAndReturn(run func
 	return _c
 }
 
-// SpendableCoins provides a mock function with given fields: ctx, addr
+// SpendableCoin provides a mock function with given fields: ctx, addr, denom
-func (_m *MockBankKeeper) SpendableCoins(ctx types.Context, addr types.AccAddress) types.Coins {
+func (_m *MockBankKeeper) SpendableCoin(ctx types.Context, addr types.AccAddress, denom string) types.Coin {
-	ret := _m.Called(ctx, addr)
+	ret := _m.Called(ctx, addr, denom)
 
 	if len(ret) == 0 {
-		panic("no return value specified for SpendableCoins")
+		panic("no return value specified for SpendableCoin")
 	}
 
-	var r0 types.Coins
+	var r0 types.Coin
-	if rf, ok := ret.Get(0).(func(types.Context, types.AccAddress) types.Coins); ok {
+	if rf, ok := ret.Get(0).(func(types.Context, types.AccAddress, string) types.Coin); ok {
-		r0 = rf(ctx, addr)
+		r0 = rf(ctx, addr, denom)
 	} else {
-		if ret.Get(0) != nil {
+		r0 = ret.Get(0).(types.Coin)
-			r0 = ret.Get(0).(types.Coins)
-		}
 	}
 
 	return r0
 }
 
-// MockBankKeeper_SpendableCoins_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'SpendableCoins'
+// MockBankKeeper_SpendableCoin_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'SpendableCoin'
-type MockBankKeeper_SpendableCoins_Call struct {
+type MockBankKeeper_SpendableCoin_Call struct {
 	*mock.Call
 }
 
-// SpendableCoins is a helper method to define mock.On call
+// SpendableCoin is a helper method to define mock.On call
 //   - ctx types.Context
 //   - addr types.AccAddress
-func (_e *MockBankKeeper_Expecter) SpendableCoins(ctx interface{}, addr interface{}) *MockBankKeeper_SpendableCoins_Call {
+//   - denom string
-	return &MockBankKeeper_SpendableCoins_Call{Call: _e.mock.On("SpendableCoins", ctx, addr)}
+func (_e *MockBankKeeper_Expecter) SpendableCoin(ctx interface{}, addr interface{}, denom interface{}) *MockBankKeeper_SpendableCoin_Call {
+	return &MockBankKeeper_SpendableCoin_Call{Call: _e.mock.On("SpendableCoin", ctx, addr, denom)}
 }
 
-func (_c *MockBankKeeper_SpendableCoins_Call) Run(run func(ctx types.Context, addr types.AccAddress)) *MockBankKeeper_SpendableCoins_Call {
+func (_c *MockBankKeeper_SpendableCoin_Call) Run(run func(ctx types.Context, addr types.AccAddress, denom string)) *MockBankKeeper_SpendableCoin_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(types.Context), args[1].(types.AccAddress))
+		run(args[0].(types.Context), args[1].(types.AccAddress), args[2].(string))
 	})
 	return _c
 }
 
-func (_c *MockBankKeeper_SpendableCoins_Call) Return(_a0 types.Coins) *MockBankKeeper_SpendableCoins_Call {
+func (_c *MockBankKeeper_SpendableCoin_Call) Return(_a0 types.Coin) *MockBankKeeper_SpendableCoin_Call {
 	_c.Call.Return(_a0)
 	return _c
 }
 
-func (_c *MockBankKeeper_SpendableCoins_Call) RunAndReturn(run func(types.Context, types.AccAddress) types.Coins) *MockBankKeeper_SpendableCoins_Call {
+func (_c *MockBankKeeper_SpendableCoin_Call) RunAndReturn(run func(types.Context, types.AccAddress, string) types.Coin) *MockBankKeeper_SpendableCoin_Call {
 	_c.Call.Return(run)
 	return _c
 }