Repository: CosmWasm/cw-multi-test
Pull request #74 (2023-09-28 14:00:09)

Extended `Wasm` trait, tested `AppBuilder::with_wasm`
# Scope of the change

## Extended `Wasm` trait with functions called on `WasmKeeper` from `App`

Providing _shortcuts_ to `Wasm` module from `App` has bounded our internal implementation of `Wasm` named `WasmKeeper` to `App`. To remove these bounds and to allow users to provide their own implementation of `Wasm` interface to `AppBuilder::with_wasm`, the `Wasm` interface has been extended with additional functions called from `App`, i.e:
- store_code,
- duplicate_code,
- contract_data,
- dump_wasm_raw

## Added integration test for `AppBuilder::with_wasm`

Added a test with example how to provide own implementation of `Wasm` interface to `AppBuilder::with_wasm` function.


Patch:
diff --git a/src/app.rs b/src/app.rs
index d1645bd..de69121 100644
--- a/src/app.rs
+++ b/src/app.rs
@@ -173,7 +173,8 @@ pub type BasicAppBuilder<ExecC, QueryC> = AppBuilder<
     FailingModule<GovMsg, Empty, Empty>,
 >;
 
-/// Utility to build App in stages. If particular items wont be set, defaults would be used
+/// Utility to build [App] in stages. If particular items/properties are not explicitly set,
+/// then default values are used.
 pub struct AppBuilder<Bank, Api, Storage, Custom, Wasm, Staking, Distr, Ibc, Gov> {
     api: Api,
     block: BlockInfo,
@@ -271,17 +272,16 @@ where
 
 impl<BankT, ApiT, StorageT, CustomT, WasmT, StakingT, DistrT, IbcT, GovT>
     AppBuilder<BankT, ApiT, StorageT, CustomT, WasmT, StakingT, DistrT, IbcT, GovT>
+where
+    CustomT: Module,
+    WasmT: Wasm<CustomT::ExecT, CustomT::QueryT>,
 {
-    /// Overwrites default wasm executor.
+    /// Overwrites the default wasm executor.
     ///
     /// At this point it is needed that new wasm implements some `Wasm` trait, but it doesn't need
     /// to be bound to Bank or Custom yet - as those may change. The cross-components validation is
     /// done on final building.
-    ///
-    /// Also it is possible to completely abandon trait bounding here which would not be bad idea,
-    /// however it might make the message on build creepy in many cases, so as for properly build
-    /// `App` we always want `Wasm` to be `Wasm`, some checks are done early.
-    pub fn with_wasm<C: Module, NewWasm: Wasm<C::ExecT, C::QueryT>>(
+    pub fn with_wasm<NewWasm: Wasm<CustomT::ExecT, CustomT::QueryT>>(
         self,
         wasm: NewWasm,
     ) -> AppBuilder<BankT, ApiT, StorageT, CustomT, NewWasm, StakingT, DistrT, IbcT, GovT> {
@@ -312,7 +312,7 @@ impl<BankT, ApiT, StorageT, CustomT, WasmT, StakingT, DistrT, IbcT, GovT>
         }
     }
 
-    /// Overwrites default bank interface
+    /// Overwrites the default bank interface.
     pub fn with_bank<NewBank: Bank>(
         self,
         bank: NewBank,
@@ -344,7 +344,7 @@ impl<BankT, ApiT, StorageT, CustomT, WasmT, StakingT, DistrT, IbcT, GovT>
         }
     }
 
-    /// Overwrites default api interface
+    /// Overwrites the default api interface.
     pub fn with_api<NewApi: Api>(
         self,
         api: NewApi,
@@ -376,7 +376,7 @@ impl<BankT, ApiT, StorageT, CustomT, WasmT, StakingT, DistrT, IbcT, GovT>
         }
     }
 
-    /// Overwrites default storage interface
+    /// Overwrites the default storage interface.
     pub fn with_storage<NewStorage: Storage>(
         self,
         storage: NewStorage,
@@ -408,15 +408,11 @@ impl<BankT, ApiT, StorageT, CustomT, WasmT, StakingT, DistrT, IbcT, GovT>
         }
     }
 
-    /// Overwrites default custom messages handler
+    /// Overwrites the default custom messages handler.
     ///
     /// At this point it is needed that new custom implements some `Module` trait, but it doesn't need
     /// to be bound to ExecC or QueryC yet - as those may change. The cross-components validation is
     /// done on final building.
-    ///
-    /// Also it is possible to completely abandon trait bounding here which would not be bad idea,
-    /// however it might make the message on build creepy in many cases, so as for properly build
-    /// `App` we always want `Wasm` to be `Wasm`, some checks are done early.
     pub fn with_custom<NewCustom: Module>(
         self,
         custom: NewCustom,
@@ -448,7 +444,7 @@ impl<BankT, ApiT, StorageT, CustomT, WasmT, StakingT, DistrT, IbcT, GovT>
         }
     }
 
-    /// Overwrites default bank interface
+    /// Overwrites the default staking interface.
     pub fn with_staking<NewStaking: Staking>(
         self,
         staking: NewStaking,
@@ -480,7 +476,7 @@ impl<BankT, ApiT, StorageT, CustomT, WasmT, StakingT, DistrT, IbcT, GovT>
         }
     }
 
-    /// Overwrites default distribution interface
+    /// Overwrites the default distribution interface.
     pub fn with_distribution<NewDistribution: Distribution>(
         self,
         distribution: NewDistribution,
@@ -513,10 +509,13 @@ impl<BankT, ApiT, StorageT, CustomT, WasmT, StakingT, DistrT, IbcT, GovT>
         }
     }
 
-    /// Overwrites default ibc interface.
+    /// Overwrites the default ibc interface.
     ///
-    /// If you wish to simply ignore/drop all returned IBC Messages, you can use the `IbcAcceptingModule` type.
-    ///     builder.with_ibc(IbcAcceptingModule::new())
+    /// If you wish to simply ignore/drop all returned IBC Messages,
+    /// you can use the `IbcAcceptingModule` type:
+    /// ```text
+    /// builder.with_ibc(IbcAcceptingModule::new())
+    /// ```
     pub fn with_ibc<NewIbc: Ibc>(
         self,
         ibc: NewIbc,
@@ -548,7 +547,7 @@ impl<BankT, ApiT, StorageT, CustomT, WasmT, StakingT, DistrT, IbcT, GovT>
         }
     }
 
-    /// Overwrites default gov interface
+    /// Overwrites the default gov interface.
     pub fn with_gov<NewGov: Gov>(
         self,
         gov: NewGov,
@@ -580,14 +579,14 @@ impl<BankT, ApiT, StorageT, CustomT, WasmT, StakingT, DistrT, IbcT, GovT>
         }
     }
 
-    /// Overwrites default initial block
+    /// Overwrites the initial block.
     pub fn with_block(mut self, block: BlockInfo) -> Self {
         self.block = block;
         self
     }
 
     /// Builds final `App`. At this point all components type have to be properly related to each
-    /// other. If there are some generics related compilation error make sure, that all components
+    /// other. If there are some generics related compilation errors, make sure that all components
     /// are properly relating to each other.
     pub fn build<F>(
         self,
@@ -668,23 +667,14 @@ where
 
 // Helper functions to call some custom WasmKeeper logic.
 // They show how we can easily add such calls to other custom keepers (CustomT, StakingT, etc)
-impl<BankT, ApiT, StorageT, CustomT, StakingT, DistrT, IbcT, GovT>
-    App<
-        BankT,
-        ApiT,
-        StorageT,
-        CustomT,
-        WasmKeeper<CustomT::ExecT, CustomT::QueryT>,
-        StakingT,
-        DistrT,
-        IbcT,
-        GovT,
-    >
+impl<BankT, ApiT, StorageT, CustomT, WasmT, StakingT, DistrT, IbcT, GovT>
+    App<BankT, ApiT, StorageT, CustomT, WasmT, StakingT, DistrT, IbcT, GovT>
 where
     BankT: Bank,
     ApiT: Api,
     StorageT: Storage,
     CustomT: Module,
+    WasmT: Wasm<CustomT::ExecT, CustomT::QueryT>,
     StakingT: Staking,
     DistrT: Distribution,
     IbcT: Ibc,
@@ -765,12 +755,12 @@ where
         self.init_modules(|router, _, _| router.wasm.duplicate_code(code_id))
     }
 
-    /// This allows to get `ContractData` for specific contract
+    /// Returns `ContractData` for the contract with specified address.
     pub fn contract_data(&self, address: &Addr) -> AnyResult<ContractData> {
-        self.read_module(|router, _, storage| router.wasm.load_contract(storage, address))
+        self.read_module(|router, _, storage| router.wasm.contract_data(storage, address))
     }
 
-    /// This gets a raw state dump of all key-values held by a given contract
+    /// Returns a raw state dump of all key-values held by a contract with specified address.
     pub fn dump_wasm_raw(&self, address: &Addr) -> Vec<Record> {
         self.read_module(|router, _, storage| router.wasm.dump_wasm_raw(storage, address))
     }
diff --git a/src/contracts.rs b/src/contracts.rs
index 76f3d6e..e362328 100644
--- a/src/contracts.rs
+++ b/src/contracts.rs
@@ -11,7 +11,7 @@ use cosmwasm_std::{
 
 use anyhow::{anyhow, bail, Result as AnyResult};
 
-/// Interface to call into a [Contract].
+/// Interface to call into a `Contract`.
 pub trait Contract<T, Q = Empty>
 where
     T: Clone + fmt::Debug + PartialEq + JsonSchema,
diff --git a/src/lib.rs b/src/lib.rs
index eb5591d..ab166c3 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -36,4 +36,4 @@ pub use crate::module::{FailingModule, Module};
 pub use crate::staking::{
     Distribution, DistributionKeeper, StakeKeeper, Staking, StakingInfo, StakingSudo,
 };
-pub use crate::wasm::{AddressGenerator, Wasm, WasmKeeper, WasmSudo};
+pub use crate::wasm::{AddressGenerator, ContractData, Wasm, WasmKeeper, WasmSudo};
diff --git a/src/module.rs b/src/module.rs
index 7125a3e..08ea118 100644
--- a/src/module.rs
+++ b/src/module.rs
@@ -8,6 +8,7 @@ use crate::AppResponse;
 use schemars::JsonSchema;
 use serde::de::DeserializeOwned;
 
+/// Interface of the module.
 pub trait Module {
     type ExecT;
     type QueryT;
diff --git a/src/wasm.rs b/src/wasm.rs
index 4cc3883..b4e3fee 100644
--- a/src/wasm.rs
+++ b/src/wasm.rs
@@ -74,8 +74,9 @@ struct CodeData {
     code_base_id: usize,
 }
 
+/// Interface to call into a `Wasm` module.
 pub trait Wasm<ExecC, QueryC> {
-    /// Handles all WasmQuery requests
+    /// Handles all `WasmQuery` requests.
     fn query(
         &self,
         api: &dyn Api,
@@ -85,7 +86,7 @@ pub trait Wasm<ExecC, QueryC> {
         request: WasmQuery,
     ) -> AnyResult<Binary>;
 
-    /// Handles all WasmMsg messages
+    /// Handles all `WasmMsg` messages.
     fn execute(
         &self,
         api: &dyn Api,
@@ -96,7 +97,7 @@ pub trait Wasm<ExecC, QueryC> {
         msg: WasmMsg,
     ) -> AnyResult<AppResponse>;
 
-    /// Admin interface, cannot be called via CosmosMsg
+    /// Handles all sudo messages, this is an admin interface and can not be called via `CosmosMsg`.
     fn sudo(
         &self,
         api: &dyn Api,
@@ -106,6 +107,19 @@ pub trait Wasm<ExecC, QueryC> {
         block: &BlockInfo,
         msg: Binary,
     ) -> AnyResult<AppResponse>;
+
+    /// Stores the contract's code and returns an identifier of the stored contract's code.
+    fn store_code(&mut self, creator: Addr, code: Box<dyn Contract<ExecC, QueryC>>) -> u64;
+
+    /// Duplicates the contract's code with specified identifier
+    /// and returns an identifier of the copy of the contract's code.
+    fn duplicate_code(&mut self, code_id: u64) -> AnyResult<u64>;
+
+    /// Returns `ContractData` for the contract with specified address.
+    fn contract_data(&self, storage: &dyn Storage, address: &Addr) -> AnyResult<ContractData>;
+
+    /// Returns a raw state dump of all key-values held by a contract with specified address.
+    fn dump_wasm_raw(&self, storage: &dyn Storage, address: &Addr) -> Vec<Record>;
 }
 
 pub struct WasmKeeper<ExecC, QueryC> {
@@ -174,7 +188,7 @@ where
             }
             WasmQuery::ContractInfo { contract_addr } => {
                 let addr = api.addr_validate(&contract_addr)?;
-                let contract = self.load_contract(storage, &addr)?;
+                let contract = self.contract_data(storage, &addr)?;
                 let mut res = ContractInfoResponse::default();
                 res.code_id = contract.code_id;
                 res.creator = contract.creator.to_string();
@@ -227,12 +241,10 @@ where
         let (res, msgs) = self.build_app_response(&contract, custom_event, res);
         self.process_response(api, router, storage, block, contract, res, msgs)
     }
-}
 
-impl<ExecC, QueryC> WasmKeeper<ExecC, QueryC> {
-    /// Stores contract code in the in-memory lookup table.
+    /// Stores the contract's code in the in-memory lookup table.
     /// Returns an identifier of the stored contract code.
-    pub fn store_code(&mut self, creator: Addr, code: Box<dyn Contract<ExecC, QueryC>>) -> u64 {
+    fn store_code(&mut self, creator: Addr, code: Box<dyn Contract<ExecC, QueryC>>) -> u64 {
         let code_base_id = self.code_base.len();
         self.code_base.push(code);
         let code_id = self.code_data.len() + 1;
@@ -244,8 +256,9 @@ impl<ExecC, QueryC> WasmKeeper<ExecC, QueryC> {
         code_id as u64
     }
 
-    /// Duplicates contract code with specified identifier.
-    pub fn duplicate_code(&mut self, code_id: u64) -> AnyResult<u64> {
+    /// Duplicates the contract's code with specified identifier.
+    /// Returns an identifier of the copy of the contract's code.
+    fn duplicate_code(&mut self, code_id: u64) -> AnyResult<u64> {
         let code_data = self.code_data(code_id)?;
         self.code_data.push(CodeData {
             creator: code_data.creator.clone(),
@@ -255,6 +268,21 @@ impl<ExecC, QueryC> WasmKeeper<ExecC, QueryC> {
         Ok(code_id + 1)
     }
 
+    /// Returns `ContractData` for the contract with specified address.
+    fn contract_data(&self, storage: &dyn Storage, address: &Addr) -> AnyResult<ContractData> {
+        CONTRACTS
+            .load(&prefixed_read(storage, NAMESPACE_WASM), address)
+            .map_err(Into::into)
+    }
+
+    /// Returns a raw state dump of all key-values held by a contract with specified address.
+    fn dump_wasm_raw(&self, storage: &dyn Storage, address: &Addr) -> Vec<Record> {
+        let storage = self.contract_storage_readonly(storage, address);
+        storage.range(None, None, Order::Ascending).collect()
+    }
+}
+
+impl<ExecC, QueryC> WasmKeeper<ExecC, QueryC> {
     /// Returns a handler to code of the contract with specified code id.
     pub fn contract_code(&self, code_id: u64) -> AnyResult<&dyn Contract<ExecC, QueryC>> {
         let code_data = self.code_data(code_id)?;
@@ -272,17 +300,6 @@ impl<ExecC, QueryC> WasmKeeper<ExecC, QueryC> {
             .ok_or(Error::UnregisteredCodeId(code_id))?)
     }
 
-    pub fn load_contract(&self, storage: &dyn Storage, address: &Addr) -> AnyResult<ContractData> {
-        CONTRACTS
-            .load(&prefixed_read(storage, NAMESPACE_WASM), address)
-            .map_err(Into::into)
-    }
-
-    pub fn dump_wasm_raw(&self, storage: &dyn Storage, address: &Addr) -> Vec<Record> {
-        let storage = self.contract_storage_readonly(storage, address);
-        storage.range(None, None, Order::Ascending).collect()
-    }
-
     fn contract_namespace(&self, contract: &Addr) -> Vec<u8> {
         let mut name = b"contract_data/".to_vec();
         name.extend_from_slice(contract.as_bytes());
@@ -433,7 +450,7 @@ where
         let admin = new_admin.map(|a| api.addr_validate(&a)).transpose()?;
 
         // check admin status
-        let mut data = self.load_contract(storage, &contract_addr)?;
+        let mut data = self.contract_data(storage, &contract_addr)?;
         if data.admin != Some(sender) {
             bail!("Only admin can update the contract admin: {:?}", data.admin);
         }
@@ -538,7 +555,7 @@ where
                 if new_code_id as usize > self.code_data.len() {
                     bail!("Cannot migrate contract to unregistered code id");
                 }
-                let mut data = self.load_contract(storage, &contract_addr)?;
+                let mut data = self.contract_data(storage, &contract_addr)?;
                 if data.admin != Some(sender) {
                     bail!("Only admin can migrate contract: {:?}", data.admin);
                 }
@@ -948,7 +965,7 @@ where
     where
         F: FnOnce(&dyn Contract<ExecC, QueryC>, Deps<QueryC>, Env) -> AnyResult<T>,
     {
-        let contract = self.load_contract(storage, &address)?;
+        let contract = self.contract_data(storage, &address)?;
         let handler = self.contract_code(contract.code_id)?;
         let storage = self.contract_storage_readonly(storage, &address);
         let env = self.get_env(address, block);
@@ -974,7 +991,7 @@ where
         F: FnOnce(&dyn Contract<ExecC, QueryC>, DepsMut<QueryC>, Env) -> AnyResult<T>,
         ExecC: DeserializeOwned,
     {
-        let contract = self.load_contract(storage, &address)?;
+        let contract = self.contract_data(storage, &address)?;
         let handler = self.contract_code(contract.code_id)?;
 
         // We don't actually need a transaction here, as it is already embedded in a transactional.
@@ -1125,7 +1142,7 @@ mod test {
 
         // verify contract data are as expected
         let contract_data = wasm_keeper
-            .load_contract(&wasm_storage, &contract_addr)
+            .contract_data(&wasm_storage, &contract_addr)
             .unwrap();
 
         assert_eq!(


Test patch:
diff --git a/tests/app_builder/mod.rs b/tests/app_builder/mod.rs
new file mode 100644
index 0000000..42eecf3
--- /dev/null
+++ b/tests/app_builder/mod.rs
@@ -0,0 +1,1 @@
+mod with_wasm;
diff --git a/tests/app_builder/with_wasm.rs b/tests/app_builder/with_wasm.rs
new file mode 100644
index 0000000..2361e0a
--- /dev/null
+++ b/tests/app_builder/with_wasm.rs
@@ -0,0 +1,126 @@
+use anyhow::Result as AnyResult;
+use cosmwasm_std::{
+    Addr, Api, Binary, BlockInfo, Deps, DepsMut, Empty, Env, MessageInfo, Querier, Record,
+    Response, StdError, Storage, SubMsg, WasmMsg, WasmQuery,
+};
+use cw_multi_test::{
+    AppBuilder, AppResponse, Contract, ContractData, ContractWrapper, CosmosRouter, Wasm,
+};
+use schemars::JsonSchema;
+use std::fmt::Debug;
+use std::marker::PhantomData;
+
+mod contracts {
+    use super::*;
+
+    pub mod caller {
+        use super::*;
+
+        fn instantiate(
+            _deps: DepsMut,
+            _env: Env,
+            _info: MessageInfo,
+            _msg: Empty,
+        ) -> Result<Response, StdError> {
+            Ok(Response::default())
+        }
+
+        fn execute(
+            _deps: DepsMut,
+            _env: Env,
+            _info: MessageInfo,
+            msg: WasmMsg,
+        ) -> Result<Response, StdError> {
+            let message = SubMsg::new(msg);
+
+            Ok(Response::new().add_submessage(message))
+        }
+
+        fn query(_deps: Deps, _env: Env, _msg: Empty) -> Result<Binary, StdError> {
+            Err(StdError::generic_err(
+                "query not implemented for the `caller` contract",
+            ))
+        }
+
+        pub fn contract<C>() -> Box<dyn Contract<C>>
+        where
+            C: Clone + Debug + PartialEq + JsonSchema + 'static,
+        {
+            let contract = ContractWrapper::new_with_empty(execute, instantiate, query);
+            Box::new(contract)
+        }
+    }
+}
+
+#[test]
+fn building_app_with_custom_wasm_should_work() {
+    struct MyWasm<ExecT, QueryT>(PhantomData<(ExecT, QueryT)>);
+
+    impl<ExecT, QueryT> Default for MyWasm<ExecT, QueryT> {
+        fn default() -> Self {
+            Self(PhantomData)
+        }
+    }
+
+    impl<ExecT, QueryT> Wasm<ExecT, QueryT> for MyWasm<ExecT, QueryT> {
+        fn query(
+            &self,
+            _api: &dyn Api,
+            _storage: &dyn Storage,
+            _querier: &dyn Querier,
+            _block: &BlockInfo,
+            _request: WasmQuery,
+        ) -> AnyResult<Binary> {
+            todo!()
+        }
+
+        fn execute(
+            &self,
+            _api: &dyn Api,
+            _storage: &mut dyn Storage,
+            _router: &dyn CosmosRouter<ExecC = ExecT, QueryC = QueryT>,
+            _block: &BlockInfo,
+            _sender: Addr,
+            _msg: WasmMsg,
+        ) -> AnyResult<AppResponse> {
+            todo!()
+        }
+
+        fn sudo(
+            &self,
+            _api: &dyn Api,
+            _contract_addr: Addr,
+            _storage: &mut dyn Storage,
+            _router: &dyn CosmosRouter<ExecC = ExecT, QueryC = QueryT>,
+            _block: &BlockInfo,
+            _msg: Binary,
+        ) -> AnyResult<AppResponse> {
+            todo!()
+        }
+
+        fn store_code(&mut self, _creator: Addr, _code: Box<dyn Contract<ExecT, QueryT>>) -> u64 {
+            154
+        }
+
+        fn duplicate_code(&mut self, _code_id: u64) -> AnyResult<u64> {
+            todo!()
+        }
+
+        fn contract_data(
+            &self,
+            _storage: &dyn Storage,
+            _address: &Addr,
+        ) -> AnyResult<ContractData> {
+            todo!()
+        }
+
+        fn dump_wasm_raw(&self, _storage: &dyn Storage, _address: &Addr) -> Vec<Record> {
+            todo!()
+        }
+    }
+
+    let app_builder = AppBuilder::default();
+    let mut app = app_builder.with_wasm(MyWasm::default()).build(|_, _, _| {});
+    let code_id = app.store_code(contracts::caller::contract());
+    assert_eq!(154, code_id);
+}
diff --git a/tests/mod.rs b/tests/mod.rs
new file mode 100644
index 0000000..a641a2e
--- /dev/null
+++ b/tests/mod.rs
@@ -0,0 +1,1 @@
+mod app_builder;