From 546403aa2f0f86b09dba7789349ba249a2d96f7f Mon Sep 17 00:00:00 2001
From: Jake Hartnell <jake.hartnell@gmail.com>
Date: Thu, 28 Dec 2023 14:27:06 -0800
Subject: [PATCH] Add staked balances query for easier frontend integration

---
 .../provider/native-staking-proxy/src/contract.rs | 13 +++++++++++--
 .../native-staking-proxy/src/multitest.rs         | 15 +++++++++++++--
 2 files changed, 24 insertions(+), 4 deletions(-)

diff --git a/contracts/provider/native-staking-proxy/src/contract.rs b/contracts/provider/native-staking-proxy/src/contract.rs
index fc7d83ef..79576013 100644
--- a/contracts/provider/native-staking-proxy/src/contract.rs
+++ b/contracts/provider/native-staking-proxy/src/contract.rs
@@ -1,7 +1,7 @@
 use cosmwasm_std::WasmMsg::Execute;
 use cosmwasm_std::{
-    coin, ensure_eq, to_binary, Coin, DistributionMsg, GovMsg, Response, StakingMsg, VoteOption,
-    WeightedVoteOption,
+    coin, ensure_eq, to_binary, Coin, Delegation, DistributionMsg, GovMsg, Response, StakingMsg,
+    VoteOption, WeightedVoteOption,
 };
 use cw2::set_contract_version;
 use cw_storage_plus::Item;
@@ -332,6 +332,15 @@ impl NativeStakingProxyContract<'_> {
     fn config(&self, ctx: QueryCtx) -> Result<ConfigResponse, ContractError> {
         Ok(self.config.load(ctx.deps.storage)?)
     }
+
+    #[msg(query)]
+    fn staked_balances(&self, ctx: QueryCtx) -> Result<Vec<Delegation>, ContractError> {
+        let balances = ctx
+            .deps
+            .querier
+            .query_all_delegations(ctx.env.contract.address)?;
+        Ok(balances)
+    }
 }
 
 // Some unit tests, due to mt limitations / unsupported msgs
diff --git a/contracts/provider/native-staking-proxy/src/multitest.rs b/contracts/provider/native-staking-proxy/src/multitest.rs
index be3d9d27..74638706 100644
--- a/contracts/provider/native-staking-proxy/src/multitest.rs
+++ b/contracts/provider/native-staking-proxy/src/multitest.rs
@@ -1,7 +1,7 @@
 use anyhow::Result as AnyResult;
 
 use cosmwasm_std::testing::mock_env;
-use cosmwasm_std::{coin, coins, to_binary, Addr, Decimal, Validator};
+use cosmwasm_std::{coin, coins, to_binary, Addr, Decimal, Delegation, Validator};
 
 use cw_multi_test::{App as MtApp, StakingInfo, StakingSudo, SudoMsg};
 
@@ -194,10 +194,21 @@ fn staking() {
     let delegation = app
         .app()
         .wrap()
-        .query_delegation(staking_proxy.contract_addr, validator.to_owned())
+        .query_delegation(staking_proxy.contract_addr.clone(), validator.to_owned())
         .unwrap()
         .unwrap();
     assert_eq!(delegation.amount, coin(120, OSMO));
+
+    // Test staked balances query
+    let balances = staking_proxy.staked_balances().unwrap();
+    assert_eq!(
+        balances,
+        vec![Delegation {
+            delegator: Addr::unchecked(proxy_addr),
+            validator: validator.to_string(),
+            amount: coin(120, OSMO),
+        }]
+    );
 }
 
 #[test]