Repository: CosmWasm/cw-multi-test
Pull request #153 (2024-03-22 08:55:55)

Expose contract storage
Added new functions to `App`: `contract_storage` and `contract_stoage_mut` that expose prefixed storage for a contract pointed by its address, with **read-only** and **read-write** access to storage respectively.


Patch:
diff --git a/src/app.rs b/src/app.rs
index b3d8db8..6608ffd 100644
--- a/src/app.rs
+++ b/src/app.rs
@@ -258,11 +258,9 @@ where
     /// Registers contract code (like uploading wasm bytecode on a chain),
     /// so it can later be used to instantiate a contract.
     pub fn store_code(&mut self, code: Box<dyn Contract<CustomT::ExecT, CustomT::QueryT>>) -> u64 {
-        self.init_modules(|router, _, _| {
-            router
-                .wasm
-                .store_code(MockApi::default().addr_make("creator"), code)
-        })
+        self.router
+            .wasm
+            .store_code(MockApi::default().addr_make("creator"), code)
     }
 
     /// Registers contract code (like [store_code](Self::store_code)),
@@ -272,7 +270,7 @@ where
         creator: Addr,
         code: Box<dyn Contract<CustomT::ExecT, CustomT::QueryT>>,
     ) -> u64 {
-        self.init_modules(|router, _, _| router.wasm.store_code(creator, code))
+        self.router.wasm.store_code(creator, code)
     }
 
     /// Registers contract code (like [store_code_with_creator](Self::store_code_with_creator)),
@@ -283,7 +281,7 @@ where
         code_id: u64,
         code: Box<dyn Contract<CustomT::ExecT, CustomT::QueryT>>,
     ) -> AnyResult<u64> {
-        self.init_modules(|router, _, _| router.wasm.store_code_with_id(creator, code_id, code))
+        self.router.wasm.store_code_with_id(creator, code_id, code)
     }
 
     /// Duplicates the contract code identified by `code_id` and returns
@@ -336,17 +334,31 @@ where
     /// assert_eq!("code id 100: no such code", app.duplicate_code(100).unwrap_err().to_string());
     /// ```
     pub fn duplicate_code(&mut self, code_id: u64) -> AnyResult<u64> {
-        self.init_modules(|router, _, _| router.wasm.duplicate_code(code_id))
+        self.router.wasm.duplicate_code(code_id)
     }
 
     /// Returns `ContractData` for the contract with specified address.
     pub fn contract_data(&self, address: &Addr) -> AnyResult<ContractData> {
-        self.read_module(|router, _, storage| router.wasm.contract_data(storage, address))
+        self.router.wasm.contract_data(&self.storage, address)
     }
 
     /// 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))
+        self.router.wasm.dump_wasm_raw(&self.storage, address)
+    }
+
+    /// Returns **read-only** storage for a contract with specified address.
+    pub fn contract_storage<'a>(&'a self, contract_addr: &Addr) -> Box<dyn Storage + 'a> {
+        self.router
+            .wasm
+            .contract_storage(&self.storage, contract_addr)
+    }
+
+    /// Returns **read-write** storage for a contract with specified address.
+    pub fn contract_storage_mut<'a>(&'a mut self, contract_addr: &Addr) -> Box<dyn Storage + 'a> {
+        self.router
+            .wasm
+            .contract_storage_mut(&mut self.storage, contract_addr)
     }
 
     /// Returns **read-only** prefixed storage with specified namespace.
diff --git a/src/wasm.rs b/src/wasm.rs
index d4644c1..35bf331 100644
--- a/src/wasm.rs
+++ b/src/wasm.rs
@@ -134,6 +134,39 @@ pub trait Wasm<ExecC, QueryC> {
 
     /// 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>;
+
+    /// Returns the namespace of the contract storage.
+    fn contract_namespace(&self, contract: &Addr) -> Vec<u8> {
+        let mut name = b"contract_data/".to_vec();
+        name.extend_from_slice(contract.as_bytes());
+        name
+    }
+
+    /// Returns **read-only** (not mutable) contract storage.
+    fn contract_storage<'a>(
+        &self,
+        storage: &'a dyn Storage,
+        address: &Addr,
+    ) -> Box<dyn Storage + 'a> {
+        // We double-namespace this, once from global storage -> wasm_storage
+        // then from wasm_storage -> the contracts subspace
+        let namespace = self.contract_namespace(address);
+        let storage = ReadonlyPrefixedStorage::multilevel(storage, &[NAMESPACE_WASM, &namespace]);
+        Box::new(storage)
+    }
+
+    /// Returns **read-write** (mutable) contract storage.
+    fn contract_storage_mut<'a>(
+        &self,
+        storage: &'a mut dyn Storage,
+        address: &Addr,
+    ) -> Box<dyn Storage + 'a> {
+        // We double-namespace this, once from global storage -> wasm_storage
+        // then from wasm_storage -> the contracts subspace
+        let namespace = self.contract_namespace(address);
+        let storage = PrefixedStorage::multilevel(storage, &[NAMESPACE_WASM, &namespace]);
+        Box::new(storage)
+    }
 }
 
 /// A structure representing a default wasm keeper.
@@ -300,7 +333,7 @@ where
 
     /// 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);
+        let storage = self.contract_storage(storage, address);
         storage.range(None, None, Order::Ascending).collect()
     }
 }
@@ -323,37 +356,6 @@ impl<ExecC, QueryC> WasmKeeper<ExecC, QueryC> {
             .ok_or_else(|| Error::unregistered_code_id(code_id))?)
     }
 
-    fn contract_namespace(&self, contract: &Addr) -> Vec<u8> {
-        let mut name = b"contract_data/".to_vec();
-        name.extend_from_slice(contract.as_bytes());
-        name
-    }
-
-    fn contract_storage<'a>(
-        &self,
-        storage: &'a mut dyn Storage,
-        address: &Addr,
-    ) -> Box<dyn Storage + 'a> {
-        // We double-namespace this, once from global storage -> wasm_storage
-        // then from wasm_storage -> the contracts subspace
-        let namespace = self.contract_namespace(address);
-        let storage = PrefixedStorage::multilevel(storage, &[NAMESPACE_WASM, &namespace]);
-        Box::new(storage)
-    }
-
-    // fails RUNTIME if you try to write. please don't
-    fn contract_storage_readonly<'a>(
-        &self,
-        storage: &'a dyn Storage,
-        address: &Addr,
-    ) -> Box<dyn Storage + 'a> {
-        // We double-namespace this, once from global storage -> wasm_storage
-        // then from wasm_storage -> the contracts subspace
-        let namespace = self.contract_namespace(address);
-        let storage = ReadonlyPrefixedStorage::multilevel(storage, &[NAMESPACE_WASM, &namespace]);
-        Box::new(storage)
-    }
-
     fn verify_attributes(attributes: &[Attribute]) -> AnyResult<()> {
         for attr in attributes {
             let key = attr.key.trim();
@@ -536,7 +538,7 @@ where
 
     /// Returns the value stored under specified key in contracts storage.
     pub fn query_raw(&self, address: Addr, storage: &dyn Storage, key: &[u8]) -> Binary {
-        let storage = self.contract_storage_readonly(storage, &address);
+        let storage = self.contract_storage(storage, &address);
         let data = storage.get(key).unwrap_or_default();
         data.into()
     }
@@ -1140,7 +1142,7 @@ where
     {
         let contract = self.contract_data(storage, &address)?;
         let handler = self.contract_code(contract.code_id)?;
-        let storage = self.contract_storage_readonly(storage, &address);
+        let storage = self.contract_storage(storage, &address);
         let env = self.get_env(address, block);
 
         let deps = Deps {
@@ -1172,7 +1174,7 @@ where
         // However, we need to get write and read access to the same storage in two different objects,
         // and this is the only way I know how to do so.
         transactional(storage, |write_cache, read_store| {
-            let mut contract_storage = self.contract_storage(write_cache, &address);
+            let mut contract_storage = self.contract_storage_mut(write_cache, &address);
             let querier = RouterQuerier::new(router, api, read_store, block);
             let env = self.get_env(address, block);
 


Test patch:
diff --git a/tests/mod.rs b/tests/mod.rs
index ddba5c2..edf9910 100644
--- a/tests/mod.rs
+++ b/tests/mod.rs
@@ -1,38 +1,36 @@
 #![cfg(test)]
 
-use cw_storage_plus::Item;
-use serde::{Deserialize, Serialize};
-
 mod test_api;
 mod test_app;
 mod test_app_builder;
+mod test_contract_storage;
 mod test_module;
 mod test_prefixed_storage;
 mod test_wasm;
 
-const COUNTER: Item<u64> = Item::new("count");
-
-#[derive(Debug, Clone, Serialize, Deserialize)]
-#[serde(rename_all = "snake_case")]
-enum CounterQueryMsg {
-    Counter {},
-}
-
-#[derive(Debug, Clone, Serialize, Deserialize)]
-struct CounterResponseMsg {
-    value: u64,
-}
-
 mod test_contracts {
-    use super::*;
 
     pub mod counter {
-        use super::*;
         use cosmwasm_std::{
             to_json_binary, Binary, Deps, DepsMut, Empty, Env, MessageInfo, Response, StdError,
             WasmMsg,
         };
         use cw_multi_test::{Contract, ContractWrapper};
+        use cw_storage_plus::Item;
+        use serde::{Deserialize, Serialize};
+
+        const COUNTER: Item<u64> = Item::new("counter");
+
+        #[derive(Debug, Clone, Serialize, Deserialize)]
+        #[serde(rename_all = "snake_case")]
+        pub enum CounterQueryMsg {
+            Counter {},
+        }
+
+        #[derive(Debug, Clone, Serialize, Deserialize)]
+        pub struct CounterResponseMsg {
+            pub value: u64,
+        }
 
         fn instantiate(
             deps: DepsMut,
diff --git a/tests/test_app_builder/test_with_storage.rs b/tests/test_app_builder/test_with_storage.rs
index 72f8998..3290f9c 100644
--- a/tests/test_app_builder/test_with_storage.rs
+++ b/tests/test_app_builder/test_with_storage.rs
@@ -1,4 +1,5 @@
-use crate::{test_contracts, CounterQueryMsg, CounterResponseMsg};
+use crate::test_contracts;
+use crate::test_contracts::counter::{CounterQueryMsg, CounterResponseMsg};
 use cosmwasm_std::{to_json_binary, Empty, Order, Record, Storage, WasmMsg};
 use cw_multi_test::{no_init, AppBuilder, Executor};
 use std::collections::BTreeMap;
diff --git a/tests/test_contract_storage/mod.rs b/tests/test_contract_storage/mod.rs
new file mode 100644
index 0000000..1ac55bb
--- /dev/null
+++ b/tests/test_contract_storage/mod.rs
@@ -0,0 +1,76 @@
+use crate::test_contracts::counter;
+use crate::test_contracts::counter::{CounterQueryMsg, CounterResponseMsg};
+use cosmwasm_std::{to_json_binary, Empty, WasmMsg};
+use cw_multi_test::{App, Executor};
+use cw_storage_plus::Item;
+
+#[test]
+fn read_write_contract_storage_should_work() {
+    // counter value saved in contract state
+    const COUNTER: Item<u64> = Item::new("counter");
+
+    // prepare the blockchain
+    let mut app = App::default();
+
+    // store the contract code
+    let creator_addr = app.api().addr_make("creator");
+    let code_id = app.store_code_with_creator(creator_addr, counter::contract());
+    assert_eq!(1, code_id);
+
+    // instantiate a new contract
+    let owner_addr = app.api().addr_make("owner");
+    let contract_addr = app
+        .instantiate_contract(code_id, owner_addr.clone(), &Empty {}, &[], "counter", None)
+        .unwrap();
+    assert!(contract_addr.as_str().starts_with("cosmwasm1"));
+
+    // `counter` contract should return value 1 after instantiation
+    let query_res: CounterResponseMsg = app
+        .wrap()
+        .query_wasm_smart(&contract_addr, &CounterQueryMsg::Counter {})
+        .unwrap();
+    assert_eq!(1, query_res.value);
+
+    {
+        // read the counter value directly from contract storage
+        let storage = app.contract_storage(&contract_addr);
+        let value = COUNTER.load(&*storage).unwrap();
+        assert_eq!(1, value);
+    }
+
+    // execute `counter` contract - this increments a counter with one
+    let execute_msg = WasmMsg::Execute {
+        contract_addr: contract_addr.clone().into(),
+        msg: to_json_binary(&Empty {}).unwrap(),
+        funds: vec![],
+    };
+    app.execute_contract(owner_addr, contract_addr.clone(), &execute_msg, &[])
+        .unwrap();
+
+    // now the `counter` contract should return value 2
+    let query_res: CounterResponseMsg = app
+        .wrap()
+        .query_wasm_smart(&contract_addr, &CounterQueryMsg::Counter {})
+        .unwrap();
+    assert_eq!(2, query_res.value);
+
+    {
+        // read the counter value directly from contract storage
+        let storage = app.contract_storage(&contract_addr);
+        let value = COUNTER.load(&*storage).unwrap();
+        assert_eq!(2, value);
+    }
+
+    {
+        // write the counter value directly into contract storage
+        let mut storage = app.contract_storage_mut(&contract_addr);
+        COUNTER.save(&mut *storage, &100).unwrap();
+    }
+
+    // now the `counter` contract should return value 100
+    let query_res: CounterResponseMsg = app
+        .wrap()
+        .query_wasm_smart(&contract_addr, &CounterQueryMsg::Counter {})
+        .unwrap();
+    assert_eq!(100, query_res.value);
+}