Repository: CosmWasm/cw-multi-test (rust)

Issue:
implement `store_code_with_id` helper
I'm working on a migrator contract which handles migrations of different contracts from V1 to V2, the V1 code ids are known and can be static inside the contract and I need to compare those with the given contracts to make sure im talking only with our modules, but with current `store_code` I can't test it with multi-test. (as a work around I am providing both code id sets from the front-end)

It can be done in 2 ways, either create another function `store_code_with_id` that will handle it specifically, or add an optional field of id to current `store_code`.

`pub fn store_code_with_id(code, id)`

We can store the highest provided code id in `self.last_code_id`, so "original" `store_code()` will be able to work, instead of using `self.codes.len()`, it will use the `last_code_id`.
Add a test that will check if the provided code id already exists results in an error.

It will cater to a specific use cases of testing contract code id against a static code id provided inside the contract, but those cases might increase as time pass on. 

Pull request description:
Added `store_code_with_id` function to `App`
closes #22

Patch:
diff --git a/.circleci/config.yml b/.circleci/config.yml
index 52675aa..f973984 100644
--- a/.circleci/config.yml
+++ b/.circleci/config.yml
@@ -155,7 +155,7 @@ jobs:
           name: Run tests with coverage
           command: |
             mkdir -p cov
-            docker run --security-opt seccomp=unconfined -v "${PWD}:/volume" xd009642/tarpaulin \
-              sh -c "cargo tarpaulin --workspace --features backtrace,cosmwasm_1_1,cosmwasm_1_2,cosmwasm_1_3,cosmwasm_1_4 --force-clean --engine llvm --out xml --output-dir cov"
+            docker run --security-opt seccomp=unconfined -v "${PWD}:/volume" xd009642/tarpaulin:0.27.1 \
+                   sh -c "cargo tarpaulin --workspace --features backtrace,cosmwasm_1_1,cosmwasm_1_2,cosmwasm_1_3,cosmwasm_1_4 --force-clean --engine llvm --out xml --output-dir cov"
       - codecov/upload:
           file: cov/cobertura.xml
diff --git a/src/app.rs b/src/app.rs
index ad035b3..ff1925c 100644
--- a/src/app.rs
+++ b/src/app.rs
@@ -256,9 +256,8 @@ where
     /// 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(Addr::unchecked("code-creator"), code)
+            let creator = MockApi::default().addr_make("creator");
+            router.wasm.store_code(creator, code)
         })
     }
 
@@ -272,6 +271,17 @@ where
         self.init_modules(|router, _, _| router.wasm.store_code(creator, code))
     }
 
+    /// Registers contract code (like [store_code_with_creator](Self::store_code_with_creator)),
+    /// but takes the code identifier as an additional argument.
+    pub fn store_code_with_id(
+        &mut self,
+        creator: Addr,
+        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))
+    }
+
     /// Duplicates the contract code identified by `code_id` and returns
     /// the identifier of the newly created copy of the contract code.
     ///
diff --git a/src/error.rs b/src/error.rs
index 013e0ae..fce588b 100644
--- a/src/error.rs
+++ b/src/error.rs
@@ -39,6 +39,14 @@ pub enum Error {
     #[error("code id {0}: no such code")]
     UnregisteredCodeId(u64),
 
+    /// Error variant for reporting duplicated contract code identifier.
+    #[error("duplicated code id {0}")]
+    DuplicatedCodeId(u64),
+
+    /// Error variant for reporting a situation when no more contract code identifiers are available.
+    #[error("no more code identifiers available")]
+    NoMoreCodeIdAvailable,
+
     /// Error variant for reporting duplicated contract addresses.
     #[error("Contract with this address already exists: {0}")]
     DuplicatedContractAddress(String),
@@ -76,7 +84,7 @@ impl Error {
     }
 
     /// Creates an instance of the [Error](Self) for invalid contract code identifier.
-    pub fn invalid_contract_code_id() -> Self {
+    pub fn invalid_code_id() -> Self {
         Self::InvalidCodeId
     }
 
@@ -85,6 +93,16 @@ impl Error {
         Self::UnregisteredCodeId(code_id)
     }
 
+    /// Creates an instance of the [Error](Self) for duplicated contract code identifier.
+    pub fn duplicated_code_id(code_id: u64) -> Self {
+        Self::DuplicatedCodeId(code_id)
+    }
+
+    /// Creates an instance of the [Error](Self) for exhausted contract code identifiers.
+    pub fn no_more_code_id_available() -> Self {
+        Self::NoMoreCodeIdAvailable
+    }
+
     /// Creates an instance of the [Error](Self) for duplicated contract addresses.
     pub fn duplicated_contract_address(address: impl Into<String>) -> Self {
         Self::DuplicatedContractAddress(address.into())
diff --git a/src/wasm.rs b/src/wasm.rs
index 9ffce73..6bc3bf6 100644
--- a/src/wasm.rs
+++ b/src/wasm.rs
@@ -19,6 +19,7 @@ use schemars::JsonSchema;
 use serde::de::DeserializeOwned;
 use serde::{Deserialize, Serialize};
 use std::borrow::Borrow;
+use std::collections::BTreeMap;
 use std::fmt::Debug;
 
 //TODO Make `CONTRACTS` private in version 1.0 when the function AddressGenerator::next_address will be removed.
@@ -71,8 +72,8 @@ struct CodeData {
     creator: Addr,
     /// Checksum of the contract's code base.
     checksum: HexBinary,
-    /// Identifier of the code base where the contract code is stored in memory.
-    code_base_id: usize,
+    /// Identifier of the _source_ code of the contract stored in wasm keeper.
+    source_id: usize,
 }
 
 /// Acts as the interface for interacting with WebAssembly (Wasm) modules.
@@ -114,6 +115,15 @@ pub trait Wasm<ExecC, QueryC> {
     /// 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;
 
+    /// Stores the contract's code under specified identifier,
+    /// returns the same code identifier when successful.
+    fn store_code_with_id(
+        &mut self,
+        creator: Addr,
+        code_id: u64,
+        code: Box<dyn Contract<ExecC, QueryC>>,
+    ) -> AnyResult<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>;
@@ -130,7 +140,7 @@ pub struct WasmKeeper<ExecC, QueryC> {
     /// Contract codes that stand for wasm code in real-life blockchain.
     code_base: Vec<Box<dyn Contract<ExecC, QueryC>>>,
     /// Code data with code base identifier and additional attributes.
-    code_data: Vec<CodeData>,
+    code_data: BTreeMap<u64, CodeData>,
     /// Contract's address generator.
     address_generator: Box<dyn AddressGenerator>,
     /// Contract's code checksum generator.
@@ -144,7 +154,7 @@ impl<ExecC, QueryC> Default for WasmKeeper<ExecC, QueryC> {
     fn default() -> Self {
         Self {
             code_base: Vec::default(),
-            code_data: Vec::default(),
+            code_data: BTreeMap::default(),
             address_generator: Box::new(SimpleAddressGenerator),
             checksum_generator: Box::new(SimpleChecksumGenerator),
             _p: std::marker::PhantomData,
@@ -231,28 +241,45 @@ where
     /// Stores the contract's code in the in-memory lookup table.
     /// Returns an identifier of the stored contract code.
     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) as u64;
-        let checksum = self.checksum_generator.checksum(&creator, code_id);
-        self.code_data.push(CodeData {
-            creator,
-            checksum,
-            code_base_id,
-        });
-        code_id
+        let code_id = self
+            .next_code_id()
+            .unwrap_or_else(|| panic!("{}", Error::NoMoreCodeIdAvailable));
+        self.save_code(code_id, creator, code)
+    }
+
+    /// Stores the contract's code in the in-memory lookup table.
+    /// Returns an identifier of the stored contract code.
+    fn store_code_with_id(
+        &mut self,
+        creator: Addr,
+        code_id: u64,
+        code: Box<dyn Contract<ExecC, QueryC>>,
+    ) -> AnyResult<u64> {
+        // validate provided contract code identifier
+        if self.code_data.contains_key(&code_id) {
+            bail!(Error::duplicated_code_id(code_id));
+        } else if code_id == 0 {
+            bail!(Error::invalid_code_id());
+        }
+        Ok(self.save_code(code_id, creator, code))
     }
 
     /// 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(),
-            checksum: code_data.checksum.clone(),
-            code_base_id: code_data.code_base_id,
-        });
-        Ok(code_id + 1)
+        let new_code_id = self
+            .next_code_id()
+            .ok_or_else(Error::no_more_code_id_available)?;
+        self.code_data.insert(
+            new_code_id,
+            CodeData {
+                creator: code_data.creator.clone(),
+                checksum: code_data.checksum.clone(),
+                source_id: code_data.source_id,
+            },
+        );
+        Ok(new_code_id)
     }
 
     /// Returns `ContractData` for the contract with specified address.
@@ -273,17 +300,17 @@ 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)?;
-        Ok(self.code_base[code_data.code_base_id].borrow())
+        Ok(self.code_base[code_data.source_id].borrow())
     }
 
     /// Returns code data of the contract with specified code id.
     fn code_data(&self, code_id: u64) -> AnyResult<&CodeData> {
         if code_id < 1 {
-            bail!(Error::invalid_contract_code_id());
+            bail!(Error::invalid_code_id());
         }
         Ok(self
             .code_data
-            .get((code_id - 1) as usize)
+            .get(&code_id)
             .ok_or_else(|| Error::unregistered_code_id(code_id))?)
     }
 
@@ -355,6 +382,35 @@ impl<ExecC, QueryC> WasmKeeper<ExecC, QueryC> {
 
         Ok(response)
     }
+
+    fn save_code(
+        &mut self,
+        code_id: u64,
+        creator: Addr,
+        code: Box<dyn Contract<ExecC, QueryC>>,
+    ) -> u64 {
+        // prepare the next identifier for the contract 'source' code
+        let source_id = self.code_base.len();
+        // calculate the checksum of the contract 'source' code based on code_id
+        let checksum = self.checksum_generator.checksum(&creator, code_id);
+        // store the 'source' code of the contract
+        self.code_base.push(code);
+        // store the additional code attributes like creator address and checksum
+        self.code_data.insert(
+            code_id,
+            CodeData {
+                creator,
+                checksum,
+                source_id,
+            },
+        );
+        code_id
+    }
+
+    /// Returns the next code identifier.
+    fn next_code_id(&self) -> Option<u64> {
+        self.code_data.keys().last().unwrap_or(&0u64).checked_add(1)
+    }
 }
 
 impl<ExecC, QueryC> WasmKeeper<ExecC, QueryC>
@@ -1154,17 +1210,14 @@ where
     }
 }
 
-// TODO: replace with code in utils
-
 #[derive(Clone, PartialEq, Message)]
 struct InstantiateResponse {
     #[prost(string, tag = "1")]
-    pub address: ::prost::alloc::string::String,
+    pub address: String,
     #[prost(bytes, tag = "2")]
-    pub data: ::prost::alloc::vec::Vec<u8>,
+    pub data: Vec<u8>,
 }
 
-// TODO: encode helpers in utils
 fn instantiate_response(data: Option<Binary>, contact_address: &Addr) -> Binary {
     let data = data.unwrap_or_default().to_vec();
     let init_data = InstantiateResponse {
@@ -1180,7 +1233,7 @@ fn instantiate_response(data: Option<Binary>, contact_address: &Addr) -> Binary
 #[derive(Clone, PartialEq, Message)]
 struct ExecuteResponse {
     #[prost(bytes, tag = "1")]
-    pub data: ::prost::alloc::vec::Vec<u8>,
+    pub data: Vec<u8>,
 }
 
 // empty return if no data present in original
@@ -1895,7 +1948,6 @@ mod test {
         let api = MockApi::default();
         let mut wasm_keeper = wasm_keeper();
         let code_id = wasm_keeper.store_code(Addr::unchecked("creator"), payout::contract());
-        assert_eq!(1, code_id);
 
         let mut wasm_storage = MockStorage::new();
 


Test patch:
diff --git a/src/tests/test_error.rs b/src/tests/test_error.rs
index eb113e7..91292f4 100644
--- a/src/tests/test_error.rs
+++ b/src/tests/test_error.rs
@@ -35,10 +35,7 @@ fn instantiating_error_should_work() {
         })
         .to_string()
     );
-    assert_eq!(
-        "code id: invalid",
-        Error::invalid_contract_code_id().to_string()
-    );
+    assert_eq!("code id: invalid", Error::invalid_code_id().to_string());
     assert_eq!(
         "code id 53: no such code",
         Error::unregistered_code_id(53).to_string()
diff --git a/tests/test_app/mod.rs b/tests/test_app/mod.rs
index 340885b..3bcc0d9 100644
--- a/tests/test_app/mod.rs
+++ b/tests/test_app/mod.rs
@@ -1,1 +1,4 @@
 mod test_instantiate2;
+mod test_store_code;
+mod test_store_code_with_creator;
+mod test_store_code_with_id;
diff --git a/tests/test_app/test_store_code.rs b/tests/test_app/test_store_code.rs
new file mode 100644
index 0000000..76ef329
--- /dev/null
+++ b/tests/test_app/test_store_code.rs
@@ -0,0 +1,35 @@
+use crate::test_contracts::counter;
+#[cfg(feature = "cosmwasm_1_2")]
+use cosmwasm_std::testing::MockApi;
+use cw_multi_test::App;
+
+#[test]
+fn storing_code_assigns_consecutive_identifiers() {
+    // prepare the application
+    let mut app = App::default();
+
+    // storing contract's code assigns consecutive code identifiers
+    for i in 1..=10 {
+        assert_eq!(i, app.store_code(counter::contract()));
+    }
+}
+
+#[test]
+#[cfg(feature = "cosmwasm_1_2")]
+fn store_code_generates_default_address_for_creator() {
+    // prepare the application
+    let mut app = App::default();
+
+    // store contract's code
+    let code_id = app.store_code(counter::contract());
+    assert_eq!(1, code_id);
+
+    // retrieve contract code info
+    let code_info_response = app.wrap().query_wasm_code_info(code_id).unwrap();
+
+    // address of the creator should be the default one
+    assert_eq!(
+        MockApi::default().addr_make("creator").as_str(),
+        code_info_response.creator
+    );
+}
diff --git a/tests/test_app/test_store_code_with_creator.rs b/tests/test_app/test_store_code_with_creator.rs
new file mode 100644
index 0000000..9cdc590
--- /dev/null
+++ b/tests/test_app/test_store_code_with_creator.rs
@@ -0,0 +1,34 @@
+#![cfg(feature = "cosmwasm_1_2")]
+
+use crate::test_contracts::counter;
+use cw_multi_test::addons::{MockApiBech32, MockApiBech32m};
+use cw_multi_test::{no_init, AppBuilder};
+
+#[test]
+fn store_code_with_custom_creator_address_should_work() {
+    // prepare the application
+    let mut app = AppBuilder::default()
+        .with_api(MockApiBech32m::new("juno"))
+        .build(no_init);
+
+    let creator = app.api().addr_make("zeus");
+
+    // store contract's code
+    let code_id = app.store_code_with_creator(creator, counter::contract());
+    assert_eq!(1, code_id);
+
+    // retrieve contract code info
+    let code_info_response = app.wrap().query_wasm_code_info(code_id).unwrap();
+
+    // address of the creator should be the custom one in Bech32m format
+    assert_eq!(
+        MockApiBech32m::new("juno").addr_make("zeus"),
+        code_info_response.creator
+    );
+
+    // address of the creator should be the custom one but not in Bech32 format
+    assert_ne!(
+        MockApiBech32::new("juno").addr_make("zeus"),
+        code_info_response.creator
+    );
+}
diff --git a/tests/test_app/test_store_code_with_id.rs b/tests/test_app/test_store_code_with_id.rs
new file mode 100644
index 0000000..382764b
--- /dev/null
+++ b/tests/test_app/test_store_code_with_id.rs
@@ -0,0 +1,67 @@
+use crate::test_contracts::counter;
+use cw_multi_test::App;
+
+#[test]
+fn storing_code_with_custom_identifier_should_work() {
+    let mut app = App::default();
+    let creator = app.api().addr_make("prometheus");
+    assert_eq!(
+        10,
+        app.store_code_with_id(creator.clone(), 10, counter::contract())
+            .unwrap()
+    );
+    assert_eq!(
+        u64::MAX,
+        app.store_code_with_id(creator, u64::MAX, counter::contract())
+            .unwrap()
+    );
+}
+
+#[test]
+fn zero_code_id_is_not_allowed() {
+    let mut app = App::default();
+    let creator = app.api().addr_make("prometheus");
+    app.store_code_with_id(creator, 0, counter::contract())
+        .unwrap_err();
+}
+
+#[test]
+fn storing_code_with_consecutive_identifiers() {
+    let mut app = App::default();
+    let creator = app.api().addr_make("prometheus");
+    assert_eq!(
+        11,
+        app.store_code_with_id(creator, 11, counter::contract())
+            .unwrap()
+    );
+    for i in 12..=20 {
+        assert_eq!(i, app.store_code(counter::contract()));
+    }
+}
+
+#[test]
+fn storing_with_the_same_id_is_not_allowed() {
+    let mut app = App::default();
+    let creator = app.api().addr_make("prometheus");
+    let code_id = 2056;
+    assert_eq!(
+        code_id,
+        app.store_code_with_id(creator.clone(), code_id, counter::contract())
+            .unwrap()
+    );
+    app.store_code_with_id(creator, code_id, counter::contract())
+        .unwrap_err();
+}
+
+#[test]
+#[should_panic(expected = "no more code identifiers available")]
+fn no_more_identifiers_available() {
+    let mut app = App::default();
+    let creator = app.api().addr_make("prometheus");
+    assert_eq!(
+        u64::MAX,
+        app.store_code_with_id(creator, u64::MAX, counter::contract())
+            .unwrap()
+    );
+    app.store_code(counter::contract());
+}
diff --git a/tests/test_app_builder/test_with_wasm.rs b/tests/test_app_builder/test_with_wasm.rs
index 13f4b05..364a32b 100644
--- a/tests/test_app_builder/test_with_wasm.rs
+++ b/tests/test_app_builder/test_with_wasm.rs
@@ -64,6 +64,15 @@ impl<ExecT, QueryT> Wasm<ExecT, QueryT> for MyWasmKeeper {
         CODE_ID
     }
 
+    fn store_code_with_id(
+        &mut self,
+        _creator: Addr,
+        code_id: u64,
+        _code: Box<dyn Contract<ExecT, QueryT>>,
+    ) -> AnyResult<u64> {
+        Ok(code_id)
+    }
+
     fn duplicate_code(&mut self, _code_id: u64) -> AnyResult<u64> {
         bail!(DUPLICATE_CODE_MSG);
     }