Move GetBalances query to GetBalance

turbofish-org · Sep 17, 2024 · 657693a · 657693a
1 parent f478fae
commit 657693a
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 28 deletions.
diff --git a/src/contract.rs b/src/contract.rs
@@ -7,7 +7,7 @@ use cosmwasm_std::{
 use cw2::set_contract_version;
 
 use crate::error::ContractError;
-use crate::msg::{ExecuteMsg, GetBalancesResponse, InstantiateMsg, QueryMsg};
+use crate::msg::{ExecuteMsg, GetBalanceResponse, InstantiateMsg, QueryMsg};
 use crate::state::{State, STATE};
 
 const CONTRACT_NAME: &str = "crates.io:osmosis-revenue-transfer-contract";
@@ -57,15 +57,7 @@ pub mod execute {
         denom: String,
     ) -> Result<Response, ContractError> {
         let state = STATE.load(deps.storage)?;
-        let coin = match deps
-            .querier
-            .query_all_balances(env.contract.address)?
-            .into_iter()
-            .find(|c| c.denom == denom)
-        {
-            Some(coin) => coin,
-            None => return Err(ContractError::DenomNotFound(denom)),
-        };
+        let coin = deps.querier.query_balance(env.contract.address, denom)?;
 
         if coin.amount < state.min_disbursal_amount.into() {
             return Err(ContractError::InsufficientFunds {});
@@ -86,17 +78,17 @@ pub mod execute {
 #[cfg_attr(not(feature = "library"), entry_point)]
 pub fn query(deps: Deps, env: Env, msg: QueryMsg) -> StdResult<Binary> {
     match msg {
-        QueryMsg::GetBalances {} => to_json_binary(&query::balance(deps, env)?),
+        QueryMsg::GetBalance { denom } => to_json_binary(&query::balance(deps, env, denom)?),
     }
 }
 
 pub mod query {
     use super::*;
 
-    pub fn balance(deps: Deps, env: Env) -> StdResult<GetBalancesResponse> {
-        let coins = deps.querier.query_all_balances(env.contract.address)?;
+    pub fn balance(deps: Deps, env: Env, denom: String) -> StdResult<GetBalanceResponse> {
+        let coin = deps.querier.query_balance(env.contract.address, denom)?;
 
-        Ok(GetBalancesResponse { balances: coins })
+        Ok(GetBalanceResponse { balance: coin })
     }
 }
 
@@ -144,16 +136,18 @@ mod tests {
         };
         let info = message_info(&Addr::unchecked("123"), &coins(2000, "denom"));
         let _res = instantiate(deps.as_mut(), mock_env(), info.clone(), msg).unwrap();
-        let msg = QueryMsg::GetBalances {};
+        let msg = QueryMsg::GetBalance {
+            denom: "denom".to_string(),
+        };
         let res = query(deps.as_ref(), mock_env(), msg).unwrap();
 
-        let value: GetBalancesResponse = from_json(&res).unwrap();
+        let value: GetBalanceResponse = from_json(&res).unwrap();
         assert_eq!(
-            vec![Coin {
+            Coin {
                 amount: Uint128::new(2000),
                 denom: "denom".to_string(),
-            }],
-            value.balances
+            },
+            value.balance
         );
     }
 }
diff --git a/src/helpers.rs b/src/helpers.rs
@@ -6,7 +6,7 @@ use cosmwasm_std::{
     WasmQuery,
 };
 
-use crate::msg::{ExecuteMsg, GetBalancesResponse, QueryMsg};
+use crate::msg::{ExecuteMsg, GetBalanceResponse, QueryMsg};
 
 #[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq, JsonSchema)]
 pub struct CwTemplateContract(pub Addr);
@@ -26,19 +26,19 @@ impl CwTemplateContract {
         .into())
     }
 
-    pub fn balance<Q, T, CQ>(&self, querier: &Q) -> StdResult<GetBalancesResponse>
+    pub fn balance<Q, T, CQ>(&self, querier: &Q, denom: String) -> StdResult<GetBalanceResponse>
     where
         Q: Querier,
         T: Into<String>,
         CQ: CustomQuery,
     {
-        let msg = QueryMsg::GetBalances {};
+        let msg = QueryMsg::GetBalance { denom };
         let query = WasmQuery::Smart {
             contract_addr: self.addr().into(),
             msg: to_json_binary(&msg)?,
         }
         .into();
-        let res: GetBalancesResponse = QuerierWrapper::<CQ>::new(querier).query(&query)?;
+        let res: GetBalanceResponse = QuerierWrapper::<CQ>::new(querier).query(&query)?;
         Ok(res)
     }
 }
diff --git a/src/msg.rs b/src/msg.rs
@@ -18,12 +18,11 @@ pub enum ExecuteMsg {
 #[cw_serde]
 #[derive(QueryResponses)]
 pub enum QueryMsg {
-    #[returns(GetBalancesResponse)]
-    GetBalances {},
+    #[returns(GetBalanceResponse)]
+    GetBalance { denom: String },
 }
 
-// We define a custom struct for each query response
 #[cw_serde]
-pub struct GetBalancesResponse {
-    pub balances: Vec<Coin>,
+pub struct GetBalanceResponse {
+    pub balance: Coin,
 }