cw_orch_mock/
state.rs

1use cosmwasm_std::{testing::mock_env, Addr};
2use cw_orch_core::{environment::StateInterface, CwEnvError};
3
4use std::collections::HashMap;
5
6#[derive(Clone, Debug)]
7/// Mock state for testing, stores addresses and code-ids.
8pub struct MockState {
9    /// Deployed contract code ids
10    pub code_ids: HashMap<String, u64>,
11    /// Deployed contract addresses
12    pub addresses: HashMap<String, Addr>,
13    /// Chain id of the mocked chain
14    pub chain_id: String,
15}
16
17impl MockState {
18    /// Creates a new empty mock state
19    pub fn new() -> Self {
20        Self {
21            addresses: HashMap::new(),
22            code_ids: HashMap::new(),
23            chain_id: mock_env().block.chain_id,
24        }
25    }
26    /// Creates a new empty mock state
27    pub fn new_with_chain_id(chain_id: &str) -> Self {
28        Self {
29            addresses: HashMap::new(),
30            code_ids: HashMap::new(),
31            chain_id: chain_id.to_string(),
32        }
33    }
34
35    pub fn with_chain_id(mut self, chain_id: &str) -> Self {
36        self.chain_id = chain_id.to_string();
37        self
38    }
39
40    pub fn set_chain_id(&mut self, chain_id: &str) {
41        self.chain_id = chain_id.to_string();
42    }
43}
44
45impl Default for MockState {
46    fn default() -> Self {
47        Self::new()
48    }
49}
50
51impl StateInterface for MockState {
52    fn get_address(&self, contract_id: &str) -> Result<Addr, CwEnvError> {
53        self.addresses
54            .get(contract_id)
55            .ok_or_else(|| CwEnvError::AddrNotInStore(contract_id.to_owned()))
56            .map(|val| val.to_owned())
57    }
58
59    fn set_address(&mut self, contract_id: &str, address: &Addr) {
60        self.addresses
61            .insert(contract_id.to_string(), address.to_owned());
62    }
63
64    fn remove_address(&mut self, contract_id: &str) {
65        self.addresses.remove(contract_id);
66    }
67
68    /// Get the locally-saved version of the contract's version on this network
69    fn get_code_id(&self, contract_id: &str) -> Result<u64, CwEnvError> {
70        self.code_ids
71            .get(contract_id)
72            .ok_or_else(|| CwEnvError::CodeIdNotInStore(contract_id.to_owned()))
73            .map(|val| val.to_owned())
74    }
75
76    /// Set the locally-saved version of the contract's latest version on this network
77    fn set_code_id(&mut self, contract_id: &str, code_id: u64) {
78        self.code_ids.insert(contract_id.to_string(), code_id);
79    }
80
81    fn remove_code_id(&mut self, contract_id: &str) {
82        self.code_ids.remove(contract_id);
83    }
84
85    fn get_all_addresses(&self) -> Result<HashMap<String, Addr>, CwEnvError> {
86        Ok(self.addresses.clone())
87    }
88
89    fn get_all_code_ids(&self) -> Result<HashMap<String, u64>, CwEnvError> {
90        Ok(self.code_ids.clone())
91    }
92}
93
94#[cfg(test)]
95mod test {
96    use cosmwasm_std::Addr;
97    use cw_orch_core::{environment::StateInterface, CwEnvError};
98    use speculoos::prelude::*;
99
100    use super::MockState;
101
102    const CONTRACT_ID: &str = "123";
103    const CONTRACT_ADDR: &str = "cosmos123";
104
105    #[test]
106    fn mock_state() {
107        let mut mock = MockState::default();
108
109        let unchecked_address = &Addr::unchecked(CONTRACT_ADDR);
110        let code_id = 123u64;
111
112        mock.set_address(CONTRACT_ID, unchecked_address);
113        mock.set_code_id(CONTRACT_ID, code_id);
114
115        // assert we get the right address
116        let addr = mock.get_address(CONTRACT_ID).unwrap();
117        asserting!(&"address is correct for contract_id")
118            .that(unchecked_address)
119            .is_equal_to(&addr);
120
121        // assert we get the right code_id
122        let fetched_id = mock.get_code_id(CONTRACT_ID).unwrap();
123        asserting!(&"code_id is correct for contract_id")
124            .that(&fetched_id)
125            .is_equal_to(code_id);
126
127        // assert we get AddrNotInStore error
128        let missing_id = &"456";
129        let error = mock.get_address(missing_id).unwrap_err();
130        let error_msg = CwEnvError::AddrNotInStore(String::from(*missing_id)).to_string();
131        asserting!(&(format!("Asserting we get CwEnvError: {}", error_msg)))
132            .that(&error.to_string())
133            .is_equal_to(CwEnvError::AddrNotInStore(String::from(*missing_id)).to_string());
134
135        // assert we get CodeIdNotInStore error
136        let error_msg = CwEnvError::CodeIdNotInStore(String::from(*missing_id)).to_string();
137        let error = mock.get_code_id(missing_id).unwrap_err();
138        asserting!(&(format!("Asserting we get CwEnvError: {}", error_msg)))
139            .that(&error.to_string())
140            .is_equal_to(CwEnvError::CodeIdNotInStore(String::from(*missing_id)).to_string());
141
142        // validate we can get all addresses
143        let total = mock.get_all_addresses().unwrap().len();
144        asserting!(&"total addresses is one")
145            .that(&total)
146            .is_equal_to(1);
147
148        // validate we can get all code_ids
149        let total = mock.get_all_code_ids().unwrap().len();
150        asserting!(&"total code_ids is one")
151            .that(&total)
152            .is_equal_to(1)
153    }
154}
cw_orch_mock/state.rs

cw_orch_mock/
state.rs
