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

Issue:
Function register_contract in wasm module should behave more like real blockchain
```
let contract_addr = self.register_contract(
    storage,
    code_id,
    sender.clone(),
    admin.map(Addr::unchecked),
    label,
    block.height,
)?;
```
For the same sender, checksum / code id and salt, the implementation above should generate the same address and fail, just like [here](https://github.com/CosmWasm/wasmd/blob/90d43687ce0facec3cf0c912cf5b331abfdf751c/x/wasm/keeper/keeper.go#L264).

Pull request:
Improving address generation
closes #66 

---

## Key changes
- extracted the address generator logic to separate module,
- refactored and improved the interface of `AddressGenerator` trait,
- implemented generation of predictable addresses to have a collision for the same salt,
- added test cases that show how to provide custom `Ap`i and custom `AddressGenerator` that support Bech32 address encoding.

**Users may now provide their own `AddressGenerator` implementation with predictable contract addresses in any address encoding.**


Patch:
diff --git a/src/addresses.rs b/src/addresses.rs
new file mode 100644
index 0000000..3a0b166
--- /dev/null
+++ b/src/addresses.rs
@@ -0,0 +1,129 @@
+//! # Implementation of address generators
+
+use crate::error::AnyResult;
+use crate::prefixed_storage::prefixed_read;
+use crate::wasm::{CONTRACTS, NAMESPACE_WASM};
+use cosmwasm_std::{Addr, Api, CanonicalAddr, HexBinary, Order, Storage};
+
+/// Common address generator interface.
+///
+/// The default implementation of this trait generates fully predictable
+/// addresses, no matter if [contract_address](AddressGenerator::contract_address)
+/// or [predictable_contract_address](AddressGenerator::predictable_contract_address) is used,
+/// but users should not make any assumptions about the value of the generated address.
+pub trait AddressGenerator {
+    #[deprecated(
+        since = "0.18.0",
+        note = "use `contract_address` or `predictable_contract_address` instead; will be removed in version 1.0.0"
+    )]
+    fn next_address(&self, storage: &mut dyn Storage) -> Addr {
+        //TODO After removing this function in version 1.0, make `CONTRACTS` and `NAMESPACE_WASM` private in `wasm.rs`.
+        let count = CONTRACTS
+            .range_raw(
+                &prefixed_read(storage, NAMESPACE_WASM),
+                None,
+                None,
+                Order::Ascending,
+            )
+            .count();
+        Addr::unchecked(format!("contract{}", count))
+    }
+
+    /// Generates a _non-predictable_ contract address, just like the real-life chain
+    /// returns contract address after its instantiation.
+    /// Address generated by this function is returned as a result of processing
+    /// `WasmMsg::Instantiate` message.
+    ///
+    /// The default implementation generates a contract address based
+    /// on contract's instance identifier only.
+    ///
+    /// # Example
+    ///
+    /// ```
+    /// # use cosmwasm_std::testing::{MockApi, MockStorage};
+    /// # use cw_multi_test::{AddressGenerator, SimpleAddressGenerator};
+    /// # let api = MockApi::default();
+    /// # let mut storage = MockStorage::default();
+    /// struct MyAddressGenerator;
+    ///
+    /// impl AddressGenerator for MyAddressGenerator {}
+    ///
+    /// let my_address_generator = MyAddressGenerator{};
+    ///
+    /// let addr = my_address_generator.contract_address(&api, &mut storage, 100, 0).unwrap();
+    /// assert_eq!(addr.to_string(),"contract0");
+    ///
+    /// let addr = my_address_generator.contract_address(&api, &mut storage, 100, 1).unwrap();
+    /// assert_eq!(addr.to_string(),"contract1");
+    ///
+    /// let addr = my_address_generator.contract_address(&api, &mut storage, 200, 5).unwrap();
+    /// assert_eq!(addr.to_string(),"contract5");
+    ///
+    /// let addr = my_address_generator.contract_address(&api, &mut storage, 200, 6).unwrap();
+    /// assert_eq!(addr.to_string(),"contract6");
+    /// ```
+    fn contract_address(
+        &self,
+        _api: &dyn Api,
+        _storage: &mut dyn Storage,
+        _code_id: u64,
+        instance_id: u64,
+    ) -> AnyResult<Addr> {
+        Ok(Addr::unchecked(format!("contract{instance_id}")))
+    }
+
+    /// Generates a _predictable_ contract address, just like the real-life chain
+    /// returns contract address after its instantiation using `MsgInstantiateContract2` message.
+    /// Address generated by this function is returned as a result of processing
+    /// `WasmMsg::Instantiate2` message.
+    ///
+    /// The default implementation generates a contract address based on provided salt only.
+    ///
+    /// # Example
+    ///
+    /// ```
+    /// # use cosmwasm_std::Api;
+    /// # use cosmwasm_std::testing::{MockApi, MockStorage};
+    /// # use cw_multi_test::{AddressGenerator, SimpleAddressGenerator};
+    /// # let api = MockApi::default();
+    /// # let mut storage = MockStorage::default();
+    /// # let creator = api.addr_canonicalize("creator").unwrap();
+    /// struct MyAddressGenerator;
+    ///
+    /// impl AddressGenerator for MyAddressGenerator {}
+    ///
+    /// let my_address_generator = MyAddressGenerator{};
+    ///
+    /// let addr = my_address_generator.predictable_contract_address(&api, &mut storage, 100, 0, &[0], &creator, &[0]).unwrap();
+    /// assert_eq!(addr.to_string(),"contract00");
+    ///
+    /// let addr = my_address_generator.predictable_contract_address(&api, &mut storage, 100, 1, &[1], &creator, &[0]).unwrap();
+    /// assert_eq!(addr.to_string(),"contract00");
+    ///
+    /// let addr = my_address_generator.predictable_contract_address(&api, &mut storage, 200, 0, &[2], &creator, &[1]).unwrap();
+    /// assert_eq!(addr.to_string(),"contract01");
+    ///
+    /// let addr = my_address_generator.predictable_contract_address(&api, &mut storage, 200, 1, &[3], &creator, &[1]).unwrap();
+    /// assert_eq!(addr.to_string(),"contract01");
+    /// ```
+    fn predictable_contract_address(
+        &self,
+        _api: &dyn Api,
+        _storage: &mut dyn Storage,
+        _code_id: u64,
+        _instance_id: u64,
+        _checksum: &[u8],
+        _creator: &CanonicalAddr,
+        salt: &[u8],
+    ) -> AnyResult<Addr> {
+        Ok(Addr::unchecked(format!(
+            "contract{}",
+            HexBinary::from(salt).to_hex()
+        )))
+    }
+}
+
+/// Default contract address generator used in [WasmKeeper](crate::WasmKeeper).
+pub struct SimpleAddressGenerator;
+
+impl AddressGenerator for SimpleAddressGenerator {}
diff --git a/src/error.rs b/src/error.rs
index 82d24af..8cd06ee 100644
--- a/src/error.rs
+++ b/src/error.rs
@@ -27,6 +27,9 @@ pub enum Error {
 
     #[error("code id {0}: no such code")]
     UnregisteredCodeId(u64),
+
+    #[error("Contract with this address already exists: {0}")]
+    DuplicatedContractAddress(String),
 }
 
 impl Error {
@@ -47,4 +50,8 @@ impl Error {
     pub fn event_type_too_short(ty: impl Into<String>) -> Self {
         Self::EventTypeTooShort(ty.into())
     }
+
+    pub fn duplicated_contract_address(address: impl Into<String>) -> Self {
+        Self::DuplicatedContractAddress(address.into())
+    }
 }
diff --git a/src/executor.rs b/src/executor.rs
index 9f216a5..90233fd 100644
--- a/src/executor.rs
+++ b/src/executor.rs
@@ -92,6 +92,40 @@ where
         Ok(Addr::unchecked(data.contract_address))
     }
 
+    /// Instantiates a new contract and returns its predictable address.
+    /// This is a helper function around [execute][Self::execute] function
+    /// with `WasmMsg::Instantiate2` message.
+    #[cfg(feature = "cosmwasm_1_2")]
+    fn instantiate2_contract<M, L, A, S>(
+        &mut self,
+        code_id: u64,
+        sender: Addr,
+        init_msg: &M,
+        funds: &[Coin],
+        label: L,
+        admin: A,
+        salt: S,
+    ) -> AnyResult<Addr>
+    where
+        M: Serialize,
+        L: Into<String>,
+        A: Into<Option<String>>,
+        S: Into<Binary>,
+    {
+        let msg = WasmMsg::Instantiate2 {
+            admin: admin.into(),
+            code_id,
+            msg: to_binary(init_msg)?,
+            funds: funds.to_vec(),
+            label: label.into(),
+            salt: salt.into(),
+        };
+        let execute_response = self.execute(sender, msg.into())?;
+        let instantiate_response =
+            parse_instantiate_response_data(execute_response.data.unwrap_or_default().as_slice())?;
+        Ok(Addr::unchecked(instantiate_response.contract_address))
+    }
+
     /// Execute a contract and process all returned messages.
     /// This is just a helper around execute(),
     /// but we parse out the data field to that what is returned by the contract (not the protobuf wrapper)
diff --git a/src/lib.rs b/src/lib.rs
index 2064e24..2ebde75 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -6,6 +6,7 @@
 //!
 //! To understand the design of this module, please refer to `../DESIGN.md`
 
+mod addresses;
 mod app;
 mod app_builder;
 mod bank;
@@ -25,6 +26,7 @@ mod tests;
 mod transactions;
 mod wasm;
 
+pub use crate::addresses::{AddressGenerator, SimpleAddressGenerator};
 pub use crate::app::{custom_app, next_block, App, BasicApp, CosmosRouter, Router, SudoMsg};
 pub use crate::app_builder::{AppBuilder, BasicAppBuilder};
 pub use crate::bank::{Bank, BankKeeper, BankSudo};
@@ -37,4 +39,4 @@ pub use crate::module::{FailingModule, Module};
 pub use crate::staking::{
     Distribution, DistributionKeeper, StakeKeeper, Staking, StakingInfo, StakingSudo,
 };
-pub use crate::wasm::{AddressGenerator, ContractData, Wasm, WasmKeeper, WasmSudo};
+pub use crate::wasm::{ContractData, Wasm, WasmKeeper, WasmSudo};
diff --git a/src/wasm.rs b/src/wasm.rs
index a9a1754..01c4885 100644
--- a/src/wasm.rs
+++ b/src/wasm.rs
@@ -1,3 +1,4 @@
+use crate::addresses::{AddressGenerator, SimpleAddressGenerator};
 use crate::app::{CosmosRouter, RouterQuerier};
 use crate::checksums::{ChecksumGenerator, SimpleChecksumGenerator};
 use crate::contracts::Contract;
@@ -20,10 +21,12 @@ use serde::{Deserialize, Serialize};
 use std::borrow::Borrow;
 use std::fmt::Debug;
 
+//TODO Make `CONTRACTS` private in version 1.0 when the function AddressGenerator::next_address will be removed.
 /// Contract state kept in storage, separate from the contracts themselves (contract code).
-const CONTRACTS: Map<&Addr, ContractData> = Map::new("contracts");
+pub(crate) const CONTRACTS: Map<&Addr, ContractData> = Map::new("contracts");
 
-pub const NAMESPACE_WASM: &[u8] = b"wasm";
+//TODO Make `NAMESPACE_WASM` private in version 1.0 when the function AddressGenerator::next_address will be removed.
+pub(crate) const NAMESPACE_WASM: &[u8] = b"wasm";
 /// See <https://github.com/chipshort/wasmd/blob/d0e3ed19f041e65f112d8e800416b3230d0005a2/x/wasm/types/events.go#L58>
 const CONTRACT_ATTR: &str = "_contract_address";
 
@@ -129,33 +132,13 @@ pub struct WasmKeeper<ExecC, QueryC> {
     _p: std::marker::PhantomData<QueryC>,
 }
 
-pub trait AddressGenerator {
-    fn next_address(&self, storage: &mut dyn Storage) -> Addr;
-}
-
-struct SimpleAddressGenerator();
-
-impl AddressGenerator for SimpleAddressGenerator {
-    fn next_address(&self, storage: &mut dyn Storage) -> Addr {
-        let count = CONTRACTS
-            .range_raw(
-                &prefixed_read(storage, NAMESPACE_WASM),
-                None,
-                None,
-                Order::Ascending,
-            )
-            .count();
-        Addr::unchecked(format!("contract{}", count))
-    }
-}
-
 impl<ExecC, QueryC> Default for WasmKeeper<ExecC, QueryC> {
     /// Returns the default value for [WasmKeeper].
     fn default() -> Self {
         Self {
             code_base: Vec::default(),
             code_data: Vec::default(),
-            address_generator: Box::new(SimpleAddressGenerator()),
+            address_generator: Box::new(SimpleAddressGenerator),
             checksum_generator: Box::new(SimpleChecksumGenerator),
             _p: std::marker::PhantomData,
         }
@@ -376,6 +359,10 @@ where
         Self::default()
     }
 
+    #[deprecated(
+        since = "0.18.0",
+        note = "use `WasmKeeper::new().with_address_generator` instead; will be removed in version 1.0.0"
+    )]
     pub fn new_with_custom_address_generator(
         address_generator: impl AddressGenerator + 'static,
     ) -> Self {
@@ -385,6 +372,14 @@ where
         }
     }
 
+    pub fn with_address_generator(
+        mut self,
+        address_generator: impl AddressGenerator + 'static,
+    ) -> Self {
+        self.address_generator = Box::new(address_generator);
+        self
+    }
+
     pub fn with_checksum_generator(
         mut self,
         checksum_generator: impl ChecksumGenerator + 'static,
@@ -612,19 +607,21 @@ where
         msg: Binary,
         funds: Vec<Coin>,
         label: String,
-        _salt: Option<Binary>,
+        salt: Option<Binary>,
     ) -> AnyResult<AppResponse> {
         if label.is_empty() {
             bail!("Label is required on all contracts");
         }
 
         let contract_addr = self.register_contract(
+            api,
             storage,
             code_id,
             sender.clone(),
             admin.map(Addr::unchecked),
             label,
             block.height,
+            salt,
         )?;
 
         // move the cash
@@ -824,24 +821,54 @@ where
         Ok(AppResponse { events, data })
     }
 
-    /// This just creates an address and empty storage instance, returning the new address
-    /// You must call init after this to set up the contract properly.
-    /// These are separated into two steps to have cleaner return values.
+    /// Creates a contract address and empty storage instance.
+    /// Returns the new contract address.
+    ///
+    /// You have to call init after this to set up the contract properly.
+    /// These two steps are separated to have cleaner return values.
     pub fn register_contract(
         &self,
+        api: &dyn Api,
         storage: &mut dyn Storage,
         code_id: u64,
         creator: Addr,
         admin: impl Into<Option<Addr>>,
         label: String,
         created: u64,
+        salt: impl Into<Option<Binary>>,
     ) -> AnyResult<Addr> {
+        // check if the contract's code with specified code_id exists
         if code_id as usize > self.code_data.len() {
             bail!("Cannot init contract with unregistered code id");
         }
 
-        let addr = self.address_generator.next_address(storage);
+        // generate a new contract address
+        let instance_id = self.instance_count(storage) as u64;
+        let addr = if let Some(salt_binary) = salt.into() {
+            // generate predictable contract address when salt is provided
+            let code_data = self.code_data(code_id)?;
+            let canonical_addr = &api.addr_canonicalize(creator.as_ref())?;
+            self.address_generator.predictable_contract_address(
+                api,
+                storage,
+                code_id,
+                instance_id,
+                code_data.checksum.as_slice(),
+                canonical_addr,
+                salt_binary.as_slice(),
+            )?
+        } else {
+            // generate non-predictable contract address
+            self.address_generator
+                .contract_address(api, storage, code_id, instance_id)?
+        };
+
+        // contract with the same address must not already exist
+        if self.contract_data(storage, &addr).is_ok() {
+            bail!(Error::duplicated_contract_address(addr));
+        }
 
+        // prepare contract data and save new contract instance
         let info = ContractData {
             code_id,
             creator,
@@ -1029,6 +1056,18 @@ where
             .save(&mut prefixed(storage, NAMESPACE_WASM), address, contract)
             .map_err(Into::into)
     }
+
+    /// Returns the number of all contract instances.
+    fn instance_count(&self, storage: &dyn Storage) -> usize {
+        CONTRACTS
+            .range_raw(
+                &prefixed_read(storage, NAMESPACE_WASM),
+                None,
+                None,
+                Order::Ascending,
+            )
+            .count()
+    }
 }
 
 // TODO: replace with code in utils
@@ -1082,8 +1121,8 @@ mod test {
     use crate::transactions::StorageTransaction;
     use cosmwasm_std::testing::{mock_env, mock_info, MockApi, MockQuerier, MockStorage};
     use cosmwasm_std::{
-        coin, from_slice, to_vec, BankMsg, Coin, CosmosMsg, Empty, GovMsg, IbcMsg, IbcQuery,
-        StdError,
+        coin, from_slice, to_vec, BankMsg, CanonicalAddr, Coin, CosmosMsg, Empty, GovMsg,
+        HexBinary, IbcMsg, IbcQuery, StdError,
     };
 
     /// Type alias for default build `Router` to make its reference in typical scenario
@@ -1124,12 +1163,14 @@ mod test {
         transactional(&mut wasm_storage, |cache, _| {
             // cannot register contract with unregistered codeId
             wasm_keeper.register_contract(
+                &api,
                 cache,
                 code_id + 1,
                 Addr::unchecked("foobar"),
                 Addr::unchecked("admin"),
                 "label".to_owned(),
                 1000,
+                None,
             )
         })
         .unwrap_err();
@@ -1137,12 +1178,14 @@ mod test {
         let contract_addr = transactional(&mut wasm_storage, |cache, _| {
             // we can register a new instance of this code
             wasm_keeper.register_contract(
+                &api,
                 cache,
                 code_id,
                 Addr::unchecked("foobar"),
                 Addr::unchecked("admin"),
                 "label".to_owned(),
                 1000,
+                None,
             )
         })
         .unwrap();
@@ -1220,12 +1263,14 @@ mod test {
 
         let contract_addr = wasm_keeper
             .register_contract(
+                &api,
                 &mut wasm_storage,
                 code_id,
                 Addr::unchecked(creator),
                 Addr::unchecked(admin),
                 "label".to_owned(),
                 1000,
+                None,
             )
             .unwrap();
 
@@ -1323,12 +1368,14 @@ mod test {
 
         let contract_addr = wasm_keeper
             .register_contract(
+                &api,
                 &mut wasm_storage,
                 code_id,
                 Addr::unchecked("foobar"),
                 Addr::unchecked("admin"),
                 "label".to_owned(),
                 1000,
+                None,
             )
             .unwrap();
 
@@ -1375,12 +1422,14 @@ mod test {
 
         let contract_addr = wasm_keeper
             .register_contract(
+                &api,
                 &mut cache,
                 code_id,
                 Addr::unchecked("foobar"),
                 None,
                 "label".to_owned(),
                 1000,
+                None,
             )
             .unwrap();
 
@@ -1489,12 +1538,14 @@ mod test {
         let contract1 = transactional(&mut wasm_storage, |cache, _| {
             let contract = wasm_keeper
                 .register_contract(
+                    &api,
                     cache,
                     code_id,
                     Addr::unchecked("foobar"),
                     None,
                     "".to_string(),
                     1000,
+                    None,
                 )
                 .unwrap();
             let info = mock_info("foobar", &[]);
@@ -1528,12 +1579,14 @@ mod test {
             // create contract 2 and use it
             let contract2 = wasm_keeper
                 .register_contract(
+                    &api,
                     cache,
                     code_id,
                     Addr::unchecked("foobar"),
                     None,
                     "".to_owned(),
                     1000,
+                    None,
                 )
                 .unwrap();
             let info = mock_info("foobar", &[]);
@@ -1562,12 +1615,14 @@ mod test {
                 // create a contract on level 2
                 let contract3 = wasm_keeper
                     .register_contract(
+                        &api,
                         cache2,
                         code_id,
                         Addr::unchecked("foobar"),
                         None,
                         "".to_owned(),
                         1000,
+                        None,
                     )
                     .unwrap();
                 let info = mock_info("johnny", &[]);
@@ -1653,12 +1708,14 @@ mod test {
 
         let contract_addr = wasm_keeper
             .register_contract(
+                &api,
                 &mut wasm_storage,
                 code_id,
                 Addr::unchecked("creator"),
                 admin.clone(),
                 "label".to_owned(),
                 1000,
+                None,
             )
             .unwrap();
 
@@ -1752,7 +1809,8 @@ mod test {
     }
 
     #[test]
-    fn by_default_uses_simple_address_generator() {
+    fn uses_simple_address_generator_by_default() {
+        let api = MockApi::default();
         let mut wasm_keeper = wasm_keeper();
         let code_id = wasm_keeper.store_code(Addr::unchecked("creator"), payout::contract());
 
@@ -1760,36 +1818,80 @@ mod test {
 
         let contract_addr = wasm_keeper
             .register_contract(
+                &api,
                 &mut wasm_storage,
                 code_id,
                 Addr::unchecked("foobar"),
                 Addr::unchecked("admin"),
                 "label".to_owned(),
                 1000,
+                None,
             )
             .unwrap();
 
         assert_eq!(
-            "contract0", contract_addr,
+            contract_addr, "contract0",
             "default address generator returned incorrect address"
-        )
+        );
+
+        let contract_addr = wasm_keeper
+            .register_contract(
+                &api,
+                &mut wasm_storage,
+                code_id,
+                Addr::unchecked("foobar"),
+                Addr::unchecked("admin"),
+                "label".to_owned(),
+                1000,
+                Binary::from(HexBinary::from_hex("01C0FFEE").unwrap()),
+            )
+            .unwrap();
+
+        assert_eq!(
+            contract_addr, "contract01c0ffee",
+            "default address generator returned incorrect address"
+        );
     }
 
     struct TestAddressGenerator {
-        addr_to_return: Addr,
+        address: Addr,
+        predictable_address: Addr,
     }
+
     impl AddressGenerator for TestAddressGenerator {
-        fn next_address(&self, _: &mut dyn Storage) -> Addr {
-            self.addr_to_return.clone()
+        fn contract_address(
+            &self,
+            _api: &dyn Api,
+            _storage: &mut dyn Storage,
+            _code_id: u64,
+            _instance_id: u64,
+        ) -> AnyResult<Addr> {
+            Ok(self.address.clone())
+        }
+
+        fn predictable_contract_address(
+            &self,
+            _api: &dyn Api,
+            _storage: &mut dyn Storage,
+            _code_id: u64,
+            _instance_id: u64,
+            _checksum: &[u8],
+            _creator: &CanonicalAddr,
+            _salt: &[u8],
+        ) -> AnyResult<Addr> {
+            Ok(self.predictable_address.clone())
         }
     }
 
     #[test]
     fn can_use_custom_address_generator() {
-        let expected_addr = Addr::unchecked("new_test_addr");
+        let api = MockApi::default();
+        let expected_addr = Addr::unchecked("address");
+        let expected_predictable_addr = Addr::unchecked("predictable_address");
         let mut wasm_keeper: WasmKeeper<Empty, Empty> =
-            WasmKeeper::new_with_custom_address_generator(TestAddressGenerator {
-                addr_to_return: expected_addr.clone(),
+            WasmKeeper::new().with_address_generator(TestAddressGenerator {
+                address: expected_addr.clone(),
+                predictable_address: expected_predictable_addr.clone(),
             });
         let code_id = wasm_keeper.store_code(Addr::unchecked("creator"), payout::contract());
 
@@ -1797,18 +1899,54 @@ mod test {
 
         let contract_addr = wasm_keeper
             .register_contract(
+                &api,
                 &mut wasm_storage,
                 code_id,
                 Addr::unchecked("foobar"),
                 Addr::unchecked("admin"),
                 "label".to_owned(),
                 1000,
+                None,
             )
             .unwrap();
 
         assert_eq!(
-            expected_addr, contract_addr,
+            contract_addr, expected_addr,
             "custom address generator returned incorrect address"
-        )
+        );
+
+        let contract_addr = wasm_keeper
+            .register_contract(
+                &api,
+                &mut wasm_storage,
+                code_id,
+                Addr::unchecked("foobar"),
+                Addr::unchecked("admin"),
+                "label".to_owned(),
+                1000,
+                Binary::from(HexBinary::from_hex("23A74B8C").unwrap()),
+            )
+            .unwrap();
+
+        assert_eq!(
+            contract_addr, expected_predictable_addr,
+            "custom address generator returned incorrect address"
+        );
+    }
+
+    #[test]
+    #[allow(deprecated)]
+    fn remove_this_test_in_version_1_0() {
+        //TODO Remove this test in version 1.0.0 of multitest, now provided only for code coverage.
+
+        let addr_gen = TestAddressGenerator {
+            address: Addr::unchecked("a"),
+            predictable_address: Addr::unchecked("b"),
+        };
+        let mut storage = MockStorage::default();
+        let contract_addr = addr_gen.next_address(&mut storage);
+        assert_eq!(contract_addr, "contract0");
+
+        let _: WasmKeeper<Empty, Empty> = WasmKeeper::new_with_custom_address_generator(addr_gen);
     }
 }


Test patch:
diff --git a/src/tests/test_app.rs b/src/tests/test_app.rs
index 4043bac..51f7e7d 100644
--- a/src/tests/test_app.rs
+++ b/src/tests/test_app.rs
@@ -1573,7 +1573,7 @@ mod contract_instantiation {
             msg: init_msg,
             funds: vec![],
             label: "label".into(),
-            salt: [0, 1, 2, 3, 4, 5].as_slice().into(),
+            salt: [1, 2, 3, 4, 5, 6].as_slice().into(),
         };
         let res = app.execute(sender, msg.into()).unwrap();
 
@@ -1581,22 +1581,9 @@ mod contract_instantiation {
         let parsed = parse_instantiate_response_data(res.data.unwrap().as_slice()).unwrap();
         assert!(parsed.data.is_none());
 
-        // assert contract's address is exactly the predicted one
-        //
-        // REMARK:
-        //   Currently implemented address generator is used to generate
-        //   the predictable address of newly instantiated contract.
-        //
-        //   Conceptually, the address of the contract is fully predictable,
-        //   because it is just the contract0 for the first instance,
-        //   contract1 for the second and so forth.
-        //
-        //   Comparing this address to real-life blockchain and the implementation
-        //   of cosmwasm_std::instantiate2_address, this approach is totally incompatible.
-        //   This problem will be handled in the next step, please see:
-        //   https://github.com/CosmWasm/cosmwasm/issues/1873
-        //   for details.
-        assert_eq!("contract0", parsed.contract_address);
+        // assert contract's address is exactly the predicted one,
+        // in default address generator, this is like `contract` + salt in hex
+        assert_eq!(parsed.contract_address, "contract010203040506");
     }
 }
 
diff --git a/tests/mod.rs b/tests/mod.rs
index 1c997f5..cd39353 100644
--- a/tests/mod.rs
+++ b/tests/mod.rs
@@ -1,5 +1,16 @@
+#![cfg(test)]
+
+use bech32::{decode, encode, FromBase32, ToBase32, Variant};
+use cosmwasm_std::{
+    instantiate2_address, Addr, Api, CanonicalAddr, RecoverPubkeyError, StdError, StdResult,
+    Storage, VerificationError,
+};
+use cw_multi_test::error::AnyResult;
+use cw_multi_test::AddressGenerator;
 use cw_storage_plus::Item;
 use serde::{Deserialize, Serialize};
+use sha2::digest::Update;
+use sha2::{Digest, Sha256};
 
 mod test_app_builder;
 mod test_module;
@@ -64,3 +75,156 @@ mod test_contracts {
         }
     }
 }
+
+mod test_api {
+    use super::*;
+
+    pub struct MockApiBech32 {
+        prefix: &'static str,
+    }
+
+    impl MockApiBech32 {
+        pub fn new(prefix: &'static str) -> Self {
+            Self { prefix }
+        }
+    }
+
+    impl Api for MockApiBech32 {
+        fn addr_validate(&self, input: &str) -> StdResult<Addr> {
+            let canonical = self.addr_canonicalize(input)?;
+            let normalized = self.addr_humanize(&canonical)?;
+            if input != normalized {
+                Err(StdError::generic_err(
+                    "Invalid input: address not normalized",
+                ))
+            } else {
+                Ok(Addr::unchecked(input))
+            }
+        }
+
+        fn addr_canonicalize(&self, input: &str) -> StdResult<CanonicalAddr> {
+            if let Ok((prefix, decoded, Variant::Bech32)) = decode(input) {
+                if prefix == self.prefix {
+                    if let Ok(bytes) = Vec::<u8>::from_base32(&decoded) {
+                        return Ok(bytes.into());
+                    }
+                }
+            }
+            Err(StdError::generic_err("Invalid input"))
+        }
+
+        fn addr_humanize(&self, canonical: &CanonicalAddr) -> StdResult<Addr> {
+            if let Ok(encoded) = encode(
+                self.prefix,
+                canonical.as_slice().to_base32(),
+                Variant::Bech32,
+            ) {
+                Ok(Addr::unchecked(encoded))
+            } else {
+                Err(StdError::generic_err("Invalid canonical address"))
+            }
+        }
+
+        fn secp256k1_verify(
+            &self,
+            _message_hash: &[u8],
+            _signature: &[u8],
+            _public_key: &[u8],
+        ) -> Result<bool, VerificationError> {
+            unimplemented!()
+        }
+
+        fn secp256k1_recover_pubkey(
+            &self,
+            _message_hash: &[u8],
+            _signature: &[u8],
+            _recovery_param: u8,
+        ) -> Result<Vec<u8>, RecoverPubkeyError> {
+            unimplemented!()
+        }
+
+        fn ed25519_verify(
+            &self,
+            _message: &[u8],
+            _signature: &[u8],
+            _public_key: &[u8],
+        ) -> Result<bool, VerificationError> {
+            unimplemented!()
+        }
+
+        fn ed25519_batch_verify(
+            &self,
+            _messages: &[&[u8]],
+            _signatures: &[&[u8]],
+            _public_keys: &[&[u8]],
+        ) -> Result<bool, VerificationError> {
+            unimplemented!()
+        }
+
+        fn debug(&self, _message: &str) {
+            unimplemented!()
+        }
+    }
+
+    impl MockApiBech32 {
+        pub fn addr_make(&self, input: &str) -> Addr {
+            let digest = Sha256::digest(input).to_vec();
+            match encode(self.prefix, digest.to_base32(), Variant::Bech32) {
+                Ok(address) => Addr::unchecked(address),
+                Err(reason) => panic!("Generating address failed with reason: {reason}"),
+            }
+        }
+    }
+}
+
+mod test_addresses {
+    use super::*;
+
+    #[derive(Default)]
+    pub struct MockAddressGenerator;
+
+    impl AddressGenerator for MockAddressGenerator {
+        fn contract_address(
+            &self,
+            api: &dyn Api,
+            _storage: &mut dyn Storage,
+            code_id: u64,
+            instance_id: u64,
+        ) -> AnyResult<Addr> {
+            let canonical_addr = Self::instantiate_address(code_id, instance_id);
+            Ok(Addr::unchecked(api.addr_humanize(&canonical_addr)?))
+        }
+
+        fn predictable_contract_address(
+            &self,
+            api: &dyn Api,
+            _storage: &mut dyn Storage,
+            _code_id: u64,
+            _instance_id: u64,
+            checksum: &[u8],
+            creator: &CanonicalAddr,
+            salt: &[u8],
+        ) -> AnyResult<Addr> {
+            let canonical_addr = instantiate2_address(checksum, creator, salt)?;
+            Ok(Addr::unchecked(api.addr_humanize(&canonical_addr)?))
+        }
+    }
+
+    impl MockAddressGenerator {
+        // non-predictable contract address generator, see `BuildContractAddressClassic`
+        // implementation in wasmd: https://github.com/CosmWasm/wasmd/blob/main/x/wasm/keeper/addresses.go#L35-L42
+        fn instantiate_address(code_id: u64, instance_id: u64) -> CanonicalAddr {
+            let mut key = Vec::<u8>::new();
+            key.extend_from_slice(b"wasm\0");
+            key.extend_from_slice(&code_id.to_be_bytes());
+            key.extend_from_slice(&instance_id.to_be_bytes());
+            let module = Sha256::digest("module".as_bytes());
+            Sha256::new()
+                .chain(module)
+                .chain(key)
+                .finalize()
+                .to_vec()
+                .into()
+        }
+    }
+}
diff --git a/tests/test_app_builder/test_with_api.rs b/tests/test_app_builder/test_with_api.rs
index ed8a250..bf883a1 100644
--- a/tests/test_app_builder/test_with_api.rs
+++ b/tests/test_app_builder/test_with_api.rs
@@ -1,106 +1,6 @@
-use bech32::{decode, encode, FromBase32, ToBase32, Variant};
-use cosmwasm_std::{
-    Addr, Api, CanonicalAddr, HexBinary, RecoverPubkeyError, StdError, StdResult, VerificationError,
-};
+use crate::test_api::MockApiBech32;
+use cosmwasm_std::{Addr, Api, CanonicalAddr, HexBinary};
 use cw_multi_test::AppBuilder;
-use sha2::{Digest, Sha256};
-
-struct MyApi {
-    prefix: &'static str,
-}
-
-impl MyApi {
-    fn new(prefix: &'static str) -> Self {
-        Self { prefix }
-    }
-}
-
-impl Api for MyApi {
-    fn addr_validate(&self, input: &str) -> StdResult<Addr> {
-        let canonical = self.addr_canonicalize(input)?;
-        let normalized = self.addr_humanize(&canonical)?;
-        if input != normalized {
-            Err(StdError::generic_err(
-                "Invalid input: address not normalized",
-            ))
-        } else {
-            Ok(Addr::unchecked(input))
-        }
-    }
-
-    fn addr_canonicalize(&self, input: &str) -> StdResult<CanonicalAddr> {
-        if let Ok((prefix, decoded, Variant::Bech32)) = decode(input) {
-            if prefix == self.prefix {
-                if let Ok(bytes) = Vec::<u8>::from_base32(&decoded) {
-                    return Ok(bytes.into());
-                }
-            }
-        }
-        Err(StdError::generic_err("Invalid input"))
-    }
-
-    fn addr_humanize(&self, canonical: &CanonicalAddr) -> StdResult<Addr> {
-        if let Ok(encoded) = encode(
-            self.prefix,
-            canonical.as_slice().to_base32(),
-            Variant::Bech32,
-        ) {
-            Ok(Addr::unchecked(encoded))
-        } else {
-            Err(StdError::generic_err("Invalid canonical address"))
-        }
-    }
-
-    fn secp256k1_verify(
-        &self,
-        _message_hash: &[u8],
-        _signature: &[u8],
-        _public_key: &[u8],
-    ) -> Result<bool, VerificationError> {
-        unimplemented!()
-    }
-
-    fn secp256k1_recover_pubkey(
-        &self,
-        _message_hash: &[u8],
-        _signature: &[u8],
-        _recovery_param: u8,
-    ) -> Result<Vec<u8>, RecoverPubkeyError> {
-        unimplemented!()
-    }
-
-    fn ed25519_verify(
-        &self,
-        _message: &[u8],
-        _signature: &[u8],
-        _public_key: &[u8],
-    ) -> Result<bool, VerificationError> {
-        unimplemented!()
-    }
-
-    fn ed25519_batch_verify(
-        &self,
-        _messages: &[&[u8]],
-        _signatures: &[&[u8]],
-        _public_keys: &[&[u8]],
-    ) -> Result<bool, VerificationError> {
-        unimplemented!()
-    }
-
-    fn debug(&self, _message: &str) {
-        unimplemented!()
-    }
-}
-
-impl MyApi {
-    fn addr_make(&self, input: &str) -> Addr {
-        let digest = Sha256::digest(input).to_vec();
-        match encode(self.prefix, digest.to_base32(), Variant::Bech32) {
-            Ok(address) => Addr::unchecked(address),
-            Err(reason) => panic!("Generating address failed with reason: {reason}"),
-        }
-    }
-}
 
 #[test]
 fn building_app_with_custom_api_should_work() {
@@ -111,7 +11,7 @@ fn building_app_with_custom_api_should_work() {
     // create application with custom api that implements
     // Bech32 address encoding with 'juno' prefix
     let app = AppBuilder::default()
-        .with_api(MyApi::new("juno"))
+        .with_api(MockApiBech32::new("juno"))
         .build(|_, _, _| {});
 
     // check address validation function
diff --git a/tests/test_wasm/mod.rs b/tests/test_wasm/mod.rs
index e691ecf..bc053f7 100644
--- a/tests/test_wasm/mod.rs
+++ b/tests/test_wasm/mod.rs
@@ -1,1 +1,2 @@
+mod test_with_addr_gen;
 mod test_with_checksum_gen;
diff --git a/tests/test_wasm/test_with_addr_gen.rs b/tests/test_wasm/test_with_addr_gen.rs
new file mode 100644
index 0000000..2b8bd5e
--- /dev/null
+++ b/tests/test_wasm/test_with_addr_gen.rs
@@ -0,0 +1,149 @@
+use crate::test_addresses::MockAddressGenerator;
+use crate::test_api::MockApiBech32;
+use crate::test_contracts;
+use cosmwasm_std::{Addr, Empty};
+use cw_multi_test::{AppBuilder, Executor, WasmKeeper};
+
+#[test]
+fn contract_address_should_work() {
+    // prepare wasm module with custom address generator
+    let wasm_keeper: WasmKeeper<Empty, Empty> =
+        WasmKeeper::new().with_address_generator(MockAddressGenerator);
+
+    // prepare application with custom api
+    let mut app = AppBuilder::default()
+        .with_api(MockApiBech32::new("purple"))
+        .with_wasm(wasm_keeper)
+        .build(|_, _, _| {});
+
+    // store contract's code
+    let code_id = app.store_code_with_creator(
+        Addr::unchecked("creator"),
+        test_contracts::counter::contract(),
+    );
+
+    let owner = app.api().addr_make("owner");
+
+    let contract_addr_1 = app
+        .instantiate_contract(code_id, owner.clone(), &Empty {}, &[], "Counter", None)
+        .unwrap();
+
+    let contract_addr_2 = app
+        .instantiate_contract(code_id, owner, &Empty {}, &[], "Counter", None)
+        .unwrap();
+
+    // addresses of the two contract instances should be different
+    assert_ne!(contract_addr_1, contract_addr_2);
+
+    // make sure that generated addresses are in valid Bech32 encoding
+    assert_eq!(
+        contract_addr_1.to_string(),
+        "purple1mzdhwvvh22wrt07w59wxyd58822qavwkx5lcej7aqfkpqqlhaqfs5efvjk"
+    );
+    assert_eq!(
+        contract_addr_2.to_string(),
+        "purple14hj2tavq8fpesdwxxcu44rty3hh90vhujrvcmstl4zr3txmfvw9smc2vxm"
+    );
+}
+
+#[test]
+#[cfg(feature = "cosmwasm_1_2")]
+fn predictable_contract_address_should_work() {
+    // prepare wasm module with custom address generator
+    let wasm_keeper: WasmKeeper<Empty, Empty> =
+        WasmKeeper::new().with_address_generator(MockAddressGenerator);
+
+    // prepare application with custom api
+    let mut app = AppBuilder::default()
+        .with_api(MockApiBech32::new("juno"))
+        .with_wasm(wasm_keeper)
+        .build(|_, _, _| {});
+
+    let creator = app.api().addr_make("creator");
+
+    // store contract's code
+    let code_id = app.store_code_with_creator(creator.clone(), test_contracts::counter::contract());
+
+    let contract_addr_1 = app
+        .instantiate2_contract(
+            code_id,
+            creator.clone(),
+            &Empty {},
+            &[],
+            "Counter",
+            None,
+            [1, 2, 3, 4, 5, 6],
+        )
+        .unwrap();
+
+    let contract_addr_2 = app
+        .instantiate2_contract(
+            code_id,
+            creator.clone(),
+            &Empty {},
+            &[],
+            "Counter",
+            None,
+            [11, 12, 13, 14, 15, 16],
+        )
+        .unwrap();
+
+    // addresses of the two contract instances should be different
+    assert_ne!(contract_addr_1, contract_addr_2);
+
+    // instantiating a contract with the same salt should fail
+    app.instantiate2_contract(
+        code_id,
+        creator,
+        &Empty {},
+        &[],
+        "Counter",
+        None,
+        [1, 2, 3, 4, 5, 6],
+    )
+    .unwrap_err();
+}
+
+#[test]
+#[cfg(feature = "cosmwasm_1_2")]
+fn creating_contract_with_the_same_predictable_address_should_fail() {
+    // prepare wasm module with custom address generator
+    let wasm_keeper: WasmKeeper<Empty, Empty> =
+        WasmKeeper::new().with_address_generator(MockAddressGenerator);
+
+    // prepare application with custom api
+    let mut app = AppBuilder::default()
+        .with_api(MockApiBech32::new("juno"))
+        .with_wasm(wasm_keeper)
+        .build(|_, _, _| {});
+
+    let creator = app.api().addr_make("creator");
+
+    // store contract's code
+    let code_id = app.store_code_with_creator(creator.clone(), test_contracts::counter::contract());
+
+    // instantiating for the first time should work
+    app.instantiate2_contract(
+        code_id,
+        creator.clone(),
+        &Empty {},
+        &[],
+        "Counter",
+        None,
+        [1, 2, 3, 4, 5, 6],
+    )
+    .unwrap();
+
+    // this instantiating should fail,
+    // the same address is generated for a contract instance
+    app.instantiate2_contract(
+        code_id,
+        creator,
+        &Empty {},
+        &[],
+        "Counter",
+        None,
+        [1, 2, 3, 4, 5, 6],
+    )
+    .unwrap_err();
+}