Feat/improve balances query

Cargo.toml

@@ -13,15 +13,15 @@ incremental = false
 overflow-checks = true
 
 [workspace.dependencies]
-neutron-sdk = { package = "neutron-sdk", git = "https://github.com/neutron-org/neutron-sdk", branch = "feat/proposal-votes" }
+neutron-sdk = { git = "https://github.com/neutron-org/neutron-sdk", branch = "feat/improve-balances-query" }
 prost = "0.12.4"
 prost-types = "0.12.4"
 cosmos-sdk-proto = { version = "0.20.0", default-features = false }
 cosmwasm-std = { version = "1.4.1", features = [
     "stargate",
     "staking",
     "cosmwasm_1_1",
-    "cosmwasm_1_2"
+    "cosmwasm_1_2",
 ] }
 cw2 = "1.1.1"
 cw-storage-plus = "1.1.0"

contracts/neutron_interchain_queries/src/contract.rs

@@ -41,7 +41,7 @@ use neutron_sdk::interchain_queries::v047::queries::{
     query_unbonding_delegations, query_validators_signing_infos,
 };
 use neutron_sdk::interchain_queries::v047::register_queries::{
-    new_register_balance_query_msg, new_register_bank_total_supply_query_msg,
+    new_register_balances_query_msg, new_register_bank_total_supply_query_msg,
     new_register_delegator_delegations_query_msg,
     new_register_delegator_unbonding_delegations_query_msg,
     new_register_distribution_fee_pool_query_msg, new_register_gov_proposals_query_msg,
@@ -83,12 +83,12 @@ pub fn execute(
     msg: ExecuteMsg,
 ) -> NeutronResult<Response<NeutronMsg>> {
     match msg {
-        ExecuteMsg::RegisterBalanceQuery {
+        ExecuteMsg::RegisterBalancesQuery {
             connection_id,
             addr,
-            denom,
+            denoms,
             update_period,
-        } => register_balance_query(connection_id, addr, denom, update_period),
+        } => register_balances_query(connection_id, addr, denoms, update_period),
         ExecuteMsg::RegisterBankTotalSupplyQuery {
             connection_id,
             denoms,
@@ -169,13 +169,13 @@ pub fn execute(
     }
 }
 
-pub fn register_balance_query(
+pub fn register_balances_query(
     connection_id: String,
     addr: String,
-    denom: String,
+    denoms: Vec<String>,
     update_period: u64,
 ) -> NeutronResult<Response<NeutronMsg>> {
-    let msg = new_register_balance_query_msg(connection_id, addr, denom, update_period)?;
+    let msg = new_register_balances_query_msg(connection_id, addr, denoms, update_period)?;
 
     Ok(Response::new().add_message(msg))
 }

contracts/neutron_interchain_queries/src/msg.rs

@@ -9,11 +9,11 @@ pub struct InstantiateMsg {}
 #[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq, JsonSchema)]
 #[serde(rename_all = "snake_case")]
 pub enum ExecuteMsg {
-    RegisterBalanceQuery {
+    RegisterBalancesQuery {
         connection_id: String,
         update_period: u64,
         addr: String,
-        denom: String,
+        denoms: Vec<String>,
     },
     RegisterBankTotalSupplyQuery {
         connection_id: String,

contracts/neutron_interchain_queries/src/testing/tests.rs

@@ -251,20 +251,27 @@ fn build_interchain_query_gov_proposal_votes_value(proposal_id: u64) -> StorageV
     }
 }
 
-fn build_interchain_query_balance_response(addr: Addr, denom: String, amount: String) -> Binary {
+fn build_interchain_query_balances_response(addr: Addr, balances: Vec<Coin>) -> Binary {
     let converted_addr_bytes = decode_and_convert(addr.as_str()).unwrap();
 
-    let balance_key = create_account_denom_balance_key(converted_addr_bytes, denom).unwrap();
+    let s: Vec<StorageValue> = balances
+        .iter()
+        .map(|c| {
+            let balance_key =
+                create_account_denom_balance_key(converted_addr_bytes.clone(), c.denom.clone())
+                    .unwrap();
+            StorageValue {
+                storage_prefix: "".to_string(),
+                key: Binary(balance_key),
+                value: Binary(c.amount.to_string().into_bytes()),
+            }
+        })
+        .collect();
 
-    let s = StorageValue {
-        storage_prefix: "".to_string(),
-        key: Binary(balance_key),
-        value: Binary(amount.into_bytes()),
-    };
     Binary::from(
         to_string(&QueryRegisteredQueryResultResponse {
             result: InterchainQueryResult {
-                kv_results: vec![s],
+                kv_results: s,
                 height: 123456,
                 revision: 2,
             },
@@ -295,11 +302,11 @@ fn register_query(
 fn test_query_balance() {
     let mut deps = dependencies(&[]);
 
-    let msg = ExecuteMsg::RegisterBalanceQuery {
+    let msg = ExecuteMsg::RegisterBalancesQuery {
         connection_id: "connection".to_string(),
         update_period: 10,
         addr: "osmo1yz54ncxj9csp7un3xled03q6thrrhy9cztkfzs".to_string(),
-        denom: "uosmo".to_string(),
+        denoms: vec!["uosmo".to_string()],
     };
 
     let keys = register_query(&mut deps, mock_env(), mock_info("", &[]), msg);
@@ -310,10 +317,9 @@ fn test_query_balance() {
     deps.querier.add_registered_queries(1, registered_query);
     deps.querier.add_query_response(
         1,
-        build_interchain_query_balance_response(
+        build_interchain_query_balances_response(
             Addr::unchecked("osmo1yz54ncxj9csp7un3xled03q6thrrhy9cztkfzs"),
-            "uosmo".to_string(),
-            "8278104".to_string(),
+            vec![Coin::new(8278104u128, "uosmo")],
         ),
     );
     let query_balance = QueryMsg::Balance { query_id: 1 };
@@ -324,7 +330,51 @@ fn test_query_balance() {
         BalanceResponse {
             last_submitted_local_height: 987,
             balances: Balances {
-                coins: vec![Coin::new(8278104u128, "uosmo".to_string())]
+                coins: vec![Coin::new(8278104u128, "uosmo")]
+            },
+        }
+    )
+}
+
+#[test]
+fn test_query_balances() {
+    let mut deps = dependencies(&[]);
+
+    let msg = ExecuteMsg::RegisterBalancesQuery {
+        connection_id: "connection".to_string(),
+        update_period: 10,
+        addr: "osmo1yz54ncxj9csp7un3xled03q6thrrhy9cztkfzs".to_string(),
+        denoms: vec!["uosmo".to_string(), "uatom".to_string()],
+    };
+
+    let keys = register_query(&mut deps, mock_env(), mock_info("", &[]), msg);
+
+    let registered_query =
+        build_registered_query_response(1, QueryParam::Keys(keys.0), QueryType::KV, 987);
+
+    deps.querier.add_registered_queries(1, registered_query);
+    deps.querier.add_query_response(
+        1,
+        build_interchain_query_balances_response(
+            Addr::unchecked("osmo1yz54ncxj9csp7un3xled03q6thrrhy9cztkfzs"),
+            vec![
+                Coin::new(8278104u128, "uosmo"),
+                Coin::new(1234567u128, "uatom"),
+            ],
+        ),
+    );
+    let query_balance = QueryMsg::Balance { query_id: 1 };
+    let resp: BalanceResponse =
+        from_json(query(deps.as_ref(), mock_env(), query_balance).unwrap()).unwrap();
+    assert_eq!(
+        resp,
+        BalanceResponse {
+            last_submitted_local_height: 987,
+            balances: Balances {
+                coins: vec![
+                    Coin::new(8278104u128, "uosmo"),
+                    Coin::new(1234567u128, "uatom")
+                ]
             },
         }
     )

contracts/neutron_validator_test/src/contract.rs

@@ -44,7 +44,7 @@ use neutron_sdk::interchain_queries::types::{
 use neutron_sdk::interchain_queries::v045::queries::query_balance;
 use neutron_sdk::interchain_queries::v045::types::{COSMOS_SDK_TRANSFER_MSG_URL, RECIPIENT_FIELD};
 use neutron_sdk::interchain_queries::v045::{
-    new_register_balance_query_msg, new_register_transfers_query_msg,
+    new_register_balances_query_msg, new_register_transfers_query_msg,
 };
 use neutron_sdk::interchain_txs::helpers::get_port_id;
 use neutron_sdk::interchain_txs::v047::helpers::decode_acknowledgement_response;
@@ -413,7 +413,7 @@ pub fn register_balance_query(
     denom: String,
     update_period: u64,
 ) -> NeutronResult<Response<NeutronMsg>> {
-    let msg = new_register_balance_query_msg(connection_id, addr, denom, update_period)?;
+    let msg = new_register_balances_query_msg(connection_id, addr, vec![denom], update_period)?;
 
     Ok(Response::new().add_message(msg))
 }