From aa2fedaf78a9cf105158f04e87fbad7c7bc1039a Mon Sep 17 00:00:00 2001
From: Derrick Lee <derrick@dlee.dev>
Date: Tue, 11 Oct 2022 14:38:49 -0700
Subject: [PATCH] Only return bkava tokens when value_in_staked_tokens is true
 (#1348)

---
 x/earn/keeper/grpc_query.go      | 38 +++++++++++++++++++++++---------
 x/earn/keeper/grpc_query_test.go |  5 ++---
 2 files changed, 30 insertions(+), 13 deletions(-)

diff --git a/x/earn/keeper/grpc_query.go b/x/earn/keeper/grpc_query.go
index 72fd09be..0adec6bf 100644
--- a/x/earn/keeper/grpc_query.go
+++ b/x/earn/keeper/grpc_query.go
@@ -298,15 +298,21 @@ func (s queryServer) getOneAccountOneVaultDeposit(
 
 	if req.ValueInStakedTokens {
 		// Get underlying ukava amount if denom is a derivative
-		if s.keeper.liquidKeeper.IsDerivativeDenom(ctx, req.Denom) {
-			ukavaValue, err := s.keeper.liquidKeeper.GetStakedTokensForDerivatives(ctx, sdk.NewCoins(value))
-			if err != nil {
-				// This should "never" happen if IsDerivativeDenom is true
-				panic("Error getting ukava value for " + req.Denom)
-			}
-
-			value = ukavaValue
+		if !s.keeper.liquidKeeper.IsDerivativeDenom(ctx, req.Denom) {
+			return nil, status.Errorf(
+				codes.InvalidArgument,
+				"denom %s is not a derivative, ValueInStakedTokens can only be used with liquid derivatives",
+				req.Denom,
+			)
 		}
+
+		ukavaValue, err := s.keeper.liquidKeeper.GetStakedTokensForDerivatives(ctx, sdk.NewCoins(value))
+		if err != nil {
+			// This should "never" happen if IsDerivativeDenom is true
+			panic("Error getting ukava value for " + req.Denom)
+		}
+
+		value = ukavaValue
 	}
 
 	return &types.QueryDepositsResponse{
@@ -417,9 +423,8 @@ func (s queryServer) getOneAccountAllDeposits(
 		var valueInStakedTokens []sdk.Coin
 
 		for _, coin := range value {
-			// Non-bkava coins are kept as is
+			// Skip non-bkava coins
 			if !s.keeper.liquidKeeper.IsDerivativeDenom(ctx, coin.Denom) {
-				valueInStakedTokens = append(valueInStakedTokens, coin)
 				continue
 			}
 
@@ -432,7 +437,20 @@ func (s queryServer) getOneAccountAllDeposits(
 			valueInStakedTokens = append(valueInStakedTokens, ukavaValue)
 		}
 
+		var filteredShares types.VaultShares
+		for _, share := range accountShare.Shares {
+			// Remove non-bkava coins from shares as they are used to
+			// determine which value is mapped to which denom
+			// These should be in the same order as valueInStakedTokens
+			if !s.keeper.liquidKeeper.IsDerivativeDenom(ctx, share.Denom) {
+				continue
+			}
+
+			filteredShares = append(filteredShares, share)
+		}
+
 		value = valueInStakedTokens
+		accountShare.Shares = filteredShares
 	}
 
 	deposits = append(deposits, types.DepositResponse{
diff --git a/x/earn/keeper/grpc_query_test.go b/x/earn/keeper/grpc_query_test.go
index 2ed96cf9..f8e774e0 100644
--- a/x/earn/keeper/grpc_query_test.go
+++ b/x/earn/keeper/grpc_query_test.go
@@ -404,15 +404,14 @@ func (suite *grpcQueryTestSuite) TestDeposits() {
 			types.DepositResponse{
 				Depositor: acc2.String(),
 				Shares: types.VaultShares{
+					// Does not include non-bkava vaults
 					types.NewVaultShare(deposit4Amount.Denom, deposit4Amount.Amount.ToDec()),
 					types.NewVaultShare(deposit3Amount.Denom, deposit3Amount.Amount.ToDec()),
-					types.NewVaultShare(deposit1Amount.Denom, deposit1Amount.Amount.ToDec()),
 				},
 				Value: sdk.Coins{
+					// Does not include non-bkava vaults
 					sdk.NewCoin("ukava", deposit4Amount.Amount),
 					sdk.NewCoin("ukava", deposit3Amount.Amount),
-					// Deposit 1 is last as it's the same order as vault shares
-					deposit1Amount,
 				},
 			},
 			res.Deposits[0],