First impl of Mock and MockBech32

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

@@ -276,7 +276,7 @@ impl WasmQuerier for CosmWasm {
         Ok(c)
     }
 
-    fn instantiate2_addr<I: serde::Serialize + std::fmt::Debug>(
+    fn instantiate2_addr(
         &self,
         code_id: u64,
         creator: impl Into<String>,

cw-orch/src/osmosis_test_tube/core.rs

@@ -284,8 +284,7 @@ impl<S: StateInterface> TxHandler for OsmosisTestTube<S> {
 impl BankSetter for OsmosisTestTube {
     type T = OsmosisTestTubeBankQuerier;
 
-    /// It's impossible to set the balance of an address directly in OsmosisTestTub
-    /// So for this implementation, we use a weird algorithm
+    /// It's impossible to set the balance of an address directly in OsmosisTestTube
     fn set_balance(
         &mut self,
         _address: impl Into<String>,

cw-orch/src/osmosis_test_tube/queriers/wasm.rs

@@ -157,7 +157,7 @@ impl WasmQuerier for OsmosisTestTubeWasmQuerier {
         Ok(c)
     }
 
-    fn instantiate2_addr<I: serde::Serialize + std::fmt::Debug>(
+    fn instantiate2_addr(
         &self,
         code_id: u64,
         creator: impl Into<String>,

cw-orch/src/prelude.rs

@@ -31,7 +31,7 @@ pub use crate::environment::{
 };
 
 // Mock for testing
-pub use crate::mock::Mock;
+pub use crate::mock::{Mock, MockBech32};
 
 // OsmosisTestTube for testing
 #[cfg(feature = "osmosis-test-tube")]

packages/cw-orch-core/src/contract/interface_traits.rs

@@ -230,7 +230,7 @@ pub trait ConditionalUpload<Chain: CwEnv>: CwOrchUpload<Chain> {
             .wasm_querier()
             .code_id_hash(latest_uploaded_code_id)
             .map_err(Into::into)?;
-        let local_hash = chain.wasm_querier().local_hash(self)?;
+        let local_hash = Chain::Wasm::local_hash(self)?;
         Ok(local_hash == on_chain_hash)
     }
 

packages/cw-orch-core/src/environment/queriers/wasm.rs

@@ -36,13 +36,12 @@ pub trait WasmQuerier: Querier {
 
     /// Returns the checksum of the WASM file if the env supports it. Will re-upload every time if not supported.
     fn local_hash<Chain: TxHandler + QueryHandler, T: Uploadable + ContractInstance<Chain>>(
-        &self,
         contract: &T,
     ) -> Result<String, CwEnvError> {
         contract.wasm().checksum()
     }
 
-    fn instantiate2_addr<I: Serialize + std::fmt::Debug>(
+    fn instantiate2_addr(
         &self,
         code_id: u64,
         creator: impl Into<String>,

packages/cw-orch-core/src/error.rs

@@ -6,6 +6,7 @@ use std::{
     str::ParseBoolError,
 };
 
+use cosmwasm_std::Instantiate2AddressError;
 use thiserror::Error;
 
 /// cw-orchestrator error wrapper using thiserror.
@@ -29,6 +30,8 @@ pub enum CwEnvError {
     ParseIntError(#[from] ParseIntError),
     #[error(transparent)]
     ParseBoolError(#[from] ParseBoolError),
+    #[error(transparent)]
+    Instantiate2AddressError(#[from] Instantiate2AddressError),
     #[error("File must be a wasm file")]
     NotWasm,
     #[error("Could not find wasm file with name {0} in artifacts:{1} dir")]

packages/cw-orch-mock/src/bech32.rs

@@ -0,0 +1,141 @@
+use std::{cell::RefCell, rc::Rc};
+
+use cosmwasm_std::{Addr, Coin, Uint128};
+use cw_multi_test::{
+    addons::{MockAddressGenerator, MockApiBech32},
+    AppBuilder, WasmKeeper,
+};
+use cw_orch_core::{
+    environment::{BankQuerier, BankSetter, DefaultQueriers, StateInterface, TxHandler},
+    CwEnvError,
+};
+use cw_utils::NativeBalance;
+
+use crate::{queriers::bank::MockBankQuerier, MockBase, MockBech32, MockState};
+
+impl MockBase<MockApiBech32, MockState> {
+    /// Create a mock environment with the default mock state.
+    pub fn new(prefix: &'static str) -> Self {
+        MockBech32::new_custom(prefix, MockState::new())
+    }
+
+    pub fn new_with_chain_id(prefix: &'static str, chain_id: &str) -> Self {
+        let chain = MockBech32::new_custom(prefix, MockState::new());
+        chain
+            .app
+            .borrow_mut()
+            .update_block(|b| b.chain_id = chain_id.to_string());
+
+        chain
+    }
+}
+
+impl<S: StateInterface> MockBase<MockApiBech32, S> {
+    pub fn create_account(&self, account_name: impl Into<String>) -> Addr {
+        self.app.borrow().api().addr_make(&account_name.into())
+    }
+}
+
+impl Default for MockBase<MockApiBech32, MockState> {
+    fn default() -> Self {
+        MockBase::<MockApiBech32, MockState>::new_custom("mock", MockState::new())
+    }
+}
+
+impl<S: StateInterface> MockBase<MockApiBech32, S> {
+    /// Create a mock environment with a custom mock state.
+    /// The state is customizable by implementing the `StateInterface` trait on a custom struct and providing it on the custom constructor.
+    pub fn new_custom(prefix: &'static str, custom_state: S) -> Self {
+        let state = Rc::new(RefCell::new(custom_state));
+        let app = Rc::new(RefCell::new(
+            AppBuilder::new_custom()
+                .with_api(MockApiBech32::new(prefix))
+                .with_wasm(WasmKeeper::default().with_address_generator(MockAddressGenerator))
+                .build(|_, _, _| {}),
+        ));
+
+        // We create an address internally
+        let sender = app.borrow().api().addr_make("sender");
+
+        Self { sender, state, app }
+    }
+}
+
+impl<S: StateInterface> MockBech32<S> {
+    /// Set the bank balance of an address.
+    pub fn set_balance(
+        &self,
+        address: &Addr,
+        amount: Vec<cosmwasm_std::Coin>,
+    ) -> Result<(), CwEnvError> {
+        self.app
+            .borrow_mut()
+            .init_modules(|router, _, storage| router.bank.init_balance(storage, address, amount))
+            .map_err(Into::into)
+    }
+
+    /// Adds the bank balance of an address.
+    pub fn add_balance(
+        &self,
+        address: &Addr,
+        amount: Vec<cosmwasm_std::Coin>,
+    ) -> Result<(), CwEnvError> {
+        let addr = &address;
+        let b = self.query_all_balances(addr)?;
+        let new_amount = NativeBalance(b) + NativeBalance(amount);
+        self.app
+            .borrow_mut()
+            .init_modules(|router, _, storage| {
+                router
+                    .bank
+                    .init_balance(storage, addr, new_amount.into_vec())
+            })
+            .map_err(Into::into)
+    }
+
+    /// Set the balance for multiple coins at once.
+    pub fn set_balances(
+        &self,
+        balances: &[(&Addr, &[cosmwasm_std::Coin])],
+    ) -> Result<(), CwEnvError> {
+        self.app
+            .borrow_mut()
+            .init_modules(|router, _, storage| -> Result<(), CwEnvError> {
+                for (addr, coins) in balances {
+                    router.bank.init_balance(storage, addr, coins.to_vec())?;
+                }
+                Ok(())
+            })
+    }
+
+    /// Query the (bank) balance of a native token for and address.
+    /// Returns the amount of the native token.
+    pub fn query_balance(&self, address: &Addr, denom: &str) -> Result<Uint128, CwEnvError> {
+        Ok(self
+            .bank_querier()
+            .balance(address, Some(denom.to_string()))?
+            .first()
+            .map(|c| c.amount)
+            .unwrap_or_default())
+    }
+
+    /// Fetch all the balances of an address.
+    pub fn query_all_balances(
+        &self,
+        address: &Addr,
+    ) -> Result<Vec<cosmwasm_std::Coin>, CwEnvError> {
+        self.bank_querier().balance(address, None)
+    }
+}
+
+impl<S: StateInterface> BankSetter for MockBech32<S> {
+    type T = MockBankQuerier<MockApiBech32>;
+
+    fn set_balance(
+        &mut self,
+        address: impl Into<String>,
+        amount: Vec<Coin>,
+    ) -> Result<(), <Self as TxHandler>::Error> {
+        (*self).set_balance(&Addr::unchecked(address), amount)
+    }
+}