Add/bank querier

CHANGELOG.md

@@ -2,6 +2,8 @@
 
 ## Unreleased
 
+- Add `MutCwEnv` for manipulating testing environments.
+- Add `BankQuerier` as trait bound to `CwEnv`.
 - Add `WasmCodeQuerier` as trait bound to `CwEnv`.
 
 ## 0.18.2
@@ -26,7 +28,7 @@
 - Added `Deploy` to prelude
 - Add ability to provide custom state in `Deploy::set_contracts_state`
 - Breaking change: remove the `&self` dependency for the `Deploy::deployed_state_file_path` method
-- Using `dirs` instead of `shellexpand` for getting the default cw-orch state dir. 
+- Using `dirs` instead of `shellexpand` for getting the default cw-orch state dir.
 - Exposed the state_dir location
 - Added better env variable management by @Kayanski
 - Added message to enable logging if not enabled
@@ -68,13 +70,13 @@
 - Added helper to modify the chain id of cw-multi-test (Mock)
 - Added a trait to be able to commit any transaction to chain (Protobuf any type)
 - Added min gas and average gas utilization for computing the tx fee
-- Added Install Readme 
+- Added Install Readme
 - Change the state file default location for relative paths `./` --> `~/.cw-orchestrator`
 - Added env variables for customizing experience
 
 ## v0.15.0
 
-- Add `add_balance` function on the `Mock` type. 
+- Add `add_balance` function on the `Mock` type.
 
 ## v0.10.0
 

contracts/counter/tests/osmosis-test-tube.rs

@@ -46,7 +46,7 @@ fn setup<Chain: CwEnv>(chain: Chain) -> CounterContract<Chain> {
 #[test]
 fn count() {
     // Create the mock
-    let test_tube = OsmosisTestTube::new(coins(100_000_000_000, "uosmo"));
+    let mut test_tube = OsmosisTestTube::new(coins(100_000_000_000, "uosmo"));
 
     let account = test_tube
         .init_account(coins(100_000_000_000, "uosmo"))

cw-orch-daemon/src/error.rs

@@ -6,6 +6,8 @@ use thiserror::Error;
 
 #[derive(Error, Debug)]
 pub enum DaemonError {
+    #[error(transparent)]
+    CosmwasmStd(#[from] cosmwasm_std::StdError),
     #[error("Reqwest HTTP(s) Error")]
     ReqwestError(#[from] ::reqwest::Error),
     #[error("JSON Conversion Error")]

cw-orch-daemon/src/queriers.rs

@@ -52,7 +52,7 @@ mod ibc;
 mod node;
 mod staking;
 
-pub use bank::Bank;
+pub use bank::{cosmrs_to_cosmwasm_coins, Bank};
 pub use cosmwasm::CosmWasm;
 pub use feegrant::Feegrant;
 pub use ibc::Ibc;

cw-orch-daemon/src/queriers/bank.rs

@@ -1,5 +1,6 @@
 use crate::{cosmos_modules, error::DaemonError};
 use cosmrs::proto::cosmos::base::{query::v1beta1::PageRequest, v1beta1::Coin};
+use cosmwasm_std::StdError;
 use tonic::transport::Channel;
 
 use super::DaemonQuerier;
@@ -130,3 +131,10 @@ impl Bank {
         Ok(denoms_metadata.metadatas)
     }
 }
+
+pub fn cosmrs_to_cosmwasm_coins(c: Coin) -> Result<cosmwasm_std::Coin, StdError> {
+    Ok(cosmwasm_std::Coin {
+        amount: c.amount.parse()?,
+        denom: c.denom,
+    })
+}

cw-orch-daemon/src/sync/core.rs

@@ -2,15 +2,15 @@ use std::{fmt::Debug, rc::Rc, time::Duration};
 
 use super::super::{sender::Wallet, DaemonAsync};
 use crate::{
-    queriers::{DaemonQuerier, Node},
+    queriers::{cosmrs_to_cosmwasm_coins, Bank, DaemonQuerier, Node},
     CosmTxResponse, DaemonBuilder, DaemonError, DaemonState,
 };
 
 use cosmrs::tendermint::Time;
 use cosmwasm_std::{Addr, Coin};
 use cw_orch_core::{
     contract::{interface_traits::Uploadable, WasmPath},
-    environment::{ChainState, TxHandler},
+    environment::{BankQuerier, ChainState, TxHandler},
 };
 use cw_orch_traits::stargate::Stargate;
 use serde::{de::DeserializeOwned, Serialize};
@@ -204,6 +204,34 @@ impl TxHandler for Daemon {
     }
 }
 
+impl BankQuerier for Daemon {
+    fn balance(
+        &self,
+        address: impl Into<String>,
+        denom: Option<String>,
+    ) -> Result<Vec<cosmwasm_std::Coin>, <Self as TxHandler>::Error> {
+        let bank = Bank::new(self.channel());
+
+        let cosmrs_coins = self.rt_handle.block_on(bank.balance(address, denom))?;
+
+        cosmrs_coins
+            .iter()
+            .map(|c| cosmrs_to_cosmwasm_coins(c.clone()))
+            .collect::<Result<Vec<_>, _>>()
+            .map_err(Into::into)
+    }
+
+    fn supply_of(
+        &self,
+        denom: impl Into<String>,
+    ) -> Result<cosmwasm_std::Coin, <Self as TxHandler>::Error> {
+        let bank = Bank::new(self.channel());
+
+        let cosmrs_coin = self.rt_handle.block_on(bank.supply_of(denom))?;
+        cosmrs_to_cosmwasm_coins(cosmrs_coin.clone()).map_err(Into::into)
+    }
+}
+
 impl Stargate for Daemon {
     fn commit_any<R>(
         &self,

cw-orch/examples/osmosis_test_tube.rs

@@ -39,7 +39,7 @@ pub fn main() {
 // This is used for documentation only
 // This is actually only used to avoid having the `mut` keyword inside the mock_usage anchor (only necessary for set_sender)
 pub fn customize() {
-    let chain = OsmosisTestTube::new(coins(1_000_000_000_000, "uosmo"));
+    let mut chain = OsmosisTestTube::new(coins(1_000_000_000_000, "uosmo"));
 
     let mut contract_counter = CounterContract::new("mock:contract_counter", chain.clone());
 

cw-orch/src/osmosis_test_tube/core.rs

@@ -1,25 +1,21 @@
+use std::str::FromStr;
+
 use crate::contract::WasmPath;
 use crate::prelude::Uploadable;
-use cosmwasm_std::ContractInfoResponse;
-use cw_orch_core::environment::WasmCodeQuerier;
+use cosmwasm_std::{coin, Addr, ContractInfoResponse, StdResult};
+
+use cw_orch_core::environment::{BankQuerier, BankSetter, WasmCodeQuerier};
 use cw_orch_traits::stargate::Stargate;
 
 use cosmwasm_std::{Binary, BlockInfo, Coin, Timestamp, Uint128};
 use cw_multi_test::AppResponse;
-use osmosis_std::types::cosmwasm::wasm::v1::QueryCodeRequest;
-use osmosis_std::types::cosmwasm::wasm::v1::QueryCodeResponse;
-use osmosis_std::types::cosmwasm::wasm::v1::QueryContractInfoRequest;
-use osmosis_std::types::cosmwasm::wasm::v1::QueryContractInfoResponse;
-use osmosis_test_tube::Account;
-use osmosis_test_tube::Bank;
-use osmosis_test_tube::ExecuteResponse;
-use osmosis_test_tube::Gamm;
-use osmosis_test_tube::Module;
-use osmosis_test_tube::Runner;
-use osmosis_test_tube::RunnerError;
-use osmosis_test_tube::SigningAccount;
-use osmosis_test_tube::Wasm;
-use std::str::FromStr;
+use osmosis_std::types::cosmos::bank::v1beta1::{QuerySupplyOfRequest, QuerySupplyOfResponse};
+use osmosis_std::types::cosmwasm::wasm::v1::{
+    QueryCodeRequest, QueryCodeResponse, QueryContractInfoRequest, QueryContractInfoResponse,
+};
+use osmosis_test_tube::{
+    Account, Bank, ExecuteResponse, Gamm, Module, Runner, RunnerError, SigningAccount, Wasm,
+};
 
 // This should be the way to import stuff.
 // But apparently osmosis-test-tube doesn't have the same dependencies as the test-tube package
@@ -31,7 +27,6 @@ use osmosis_test_tube::osmosis_std::{
 use osmosis_test_tube::OsmosisTestApp;
 use std::{cell::RefCell, fmt::Debug, rc::Rc};
 
-use cosmwasm_std::Addr;
 use serde::{de::DeserializeOwned, Serialize};
 
 use crate::{
@@ -57,7 +52,7 @@ pub use osmosis_test_tube;
 /// use cw_orch::osmosis_test_tube::osmosis_test_tube::Account;
 ///
 /// // Creates an app, creates a sender with an initial balance
-/// let tube: OsmosisTestTube = OsmosisTestTube::new(coins(1_000_000_000_000, "uosmo"));
+/// let mut tube: OsmosisTestTube = OsmosisTestTube::new(coins(1_000_000_000_000, "uosmo"));
 ///
 /// // create an additional account
 /// let account = tube.init_account(coins(1_000_000_000, "uatom")).unwrap();
@@ -83,27 +78,33 @@ fn map_err(e: RunnerError) -> CwOrchError {
 impl<S: StateInterface> OsmosisTestTube<S> {
     /// Creates an account and sets its balance
     pub fn init_account(
-        &self,
+        &mut self,
         amount: Vec<cosmwasm_std::Coin>,
     ) -> Result<Rc<SigningAccount>, CwOrchError> {
-        self.app
+        let account = self
+            .app
             .borrow()
             .init_account(&amount)
             .map_err(map_err)
-            .map(Rc::new)
+            .map(Rc::new)?;
+
+        Ok(account)
     }
 
     /// Creates accounts and sets their balance
     pub fn init_accounts(
-        &self,
+        &mut self,
         amount: Vec<cosmwasm_std::Coin>,
         account_n: u64,
     ) -> Result<Vec<Rc<SigningAccount>>, CwOrchError> {
-        self.app
+        let accounts: Vec<_> = self
+            .app
             .borrow()
             .init_accounts(&amount, account_n)
             .map_err(map_err)
-            .map(|s| s.into_iter().map(Rc::new).collect())
+            .map(|s| s.into_iter().map(Rc::new).collect())?;
+
+        Ok(accounts)
     }
 
     /// Sends coins a specific address
@@ -151,7 +152,9 @@ impl<S: StateInterface> OsmosisTestTube<S> {
             })
             .map_err(map_err)?
             .balance
-            .map(|c| Uint128::from_str(&c.amount).unwrap())
+            .map(to_cosmwasm_coin)
+            .transpose()?
+            .map(|c| c.amount)
             .unwrap_or(Uint128::zero());
         Ok(amount)
     }
@@ -169,11 +172,8 @@ impl<S: StateInterface> OsmosisTestTube<S> {
             .map_err(map_err)?
             .balances
             .into_iter()
-            .map(|c| Coin {
-                amount: Uint128::from_str(&c.amount).unwrap(),
-                denom: c.denom,
-            })
-            .collect();
+            .map(to_cosmwasm_coin)
+            .collect::<Result<Vec<_>, _>>()?;
         Ok(amount)
     }
 }
@@ -326,6 +326,65 @@ impl<S: StateInterface> TxHandler for OsmosisTestTube<S> {
     }
 }
 
+impl BankQuerier for OsmosisTestTube {
+    fn balance(
+        &self,
+        address: impl Into<String>,
+        denom: Option<String>,
+    ) -> Result<Vec<cosmwasm_std::Coin>, <Self as TxHandler>::Error> {
+        if let Some(denom) = denom {
+            Ok(vec![Coin {
+                amount: self.query_balance(&address.into(), &denom)?,
+                denom,
+            }])
+        } else {
+            self.query_all_balances(&address.into())
+        }
+    }
+
+    fn supply_of(
+        &self,
+        denom: impl Into<String>,
+    ) -> Result<cosmwasm_std::Coin, <Self as TxHandler>::Error> {
+        let denom: String = denom.into();
+        let supply_of_result: QuerySupplyOfResponse = self
+            .app
+            .borrow()
+            .query(
+                "/cosmos.bank.v1beta1.Query/SupplyOf",
+                &QuerySupplyOfRequest {
+                    denom: denom.clone(),
+                },
+            )
+            .map_err(map_err)?;
+
+        Ok(supply_of_result
+            .amount
+            .map(|c| {
+                // Ok::<_, StdError>(cosmwasm_std::Coin {
+                //     amount: c.amount.parse()?,
+                //     denom: c.denom,
+                // })
+                to_cosmwasm_coin(c)
+            })
+            .transpose()?
+            .unwrap_or(coin(0, &denom)))
+    }
+}
+
+impl BankSetter for OsmosisTestTube {
+    /// It's impossible to set the balance of an address directly in OsmosisTestTub
+    /// So for this implementation, we use a weird algorithm
+    fn set_balance(
+        &mut self,
+        _address: &Addr,
+        _amount: Vec<Coin>,
+    ) -> Result<(), <Self as TxHandler>::Error> {
+        // We check the current balance
+        unimplemented!();
+    }
+}
+
 impl Stargate for OsmosisTestTube {
     fn commit_any<R: prost::Message + Default>(
         &self,
@@ -396,10 +455,18 @@ impl WasmCodeQuerier for OsmosisTestTube {
     }
 }
 
+fn to_cosmwasm_coin(c: osmosis_std::types::cosmos::base::v1beta1::Coin) -> StdResult<Coin> {
+    Ok(Coin {
+        amount: Uint128::from_str(&c.amount)?,
+        denom: c.denom,
+    })
+}
+
 #[cfg(test)]
 pub mod tests {
     use cosmwasm_std::{coins, ContractInfoResponse};
-    use cw_orch_core::environment::WasmCodeQuerier;
+    use cw_orch_core::environment::{BankQuerier, WasmCodeQuerier};
+
     use osmosis_test_tube::Account;
 
     use super::OsmosisTestTube;
@@ -433,4 +500,18 @@ pub mod tests {
 
         Ok(())
     }
+
+    #[test]
+    fn bank_querier_works() -> anyhow::Result<()> {
+        let denom = "urandom";
+        let init_coins = coins(45, denom);
+        let app = OsmosisTestTube::new(init_coins.clone());
+        let sender = app.sender.address();
+        assert_eq!(
+            app.balance(sender.clone(), Some(denom.to_string()))?,
+            init_coins
+        );
+        assert_eq!(app.supply_of(denom.to_string())?, init_coins[0]);
+        Ok(())
+    }
 }

packages/cw-orch-core/Cargo.toml

@@ -38,6 +38,7 @@ serde_json = { workspace = true }
 # Ethereum deps
 snailquote = { version = "0.3.1", optional = true }
 dirs = "5.0.1"
+cw-utils.workspace = true
 
 [dev-dependencies]
 speculoos = { workspace = true }