cw_orch_clone_testing/
core.rs

1use std::{cell::RefCell, fmt::Debug, io::Read, rc::Rc};
2
3use clone_cw_multi_test::{
4    addons::{MockAddressGenerator, MockApiBech32},
5    wasm_emulation::{channel::RemoteChannel, storage::analyzer::StorageAnalyzer},
6    App, AppBuilder, BankKeeper, Contract, Executor, WasmKeeper,
7};
8use cosmwasm_std::{
9    to_json_binary, Addr, BankMsg, Binary, Coin, CosmosMsg, Empty, Event, StdError, StdResult,
10    Uint128, WasmMsg,
11};
12use cw_orch_core::{
13    contract::interface_traits::{ContractInstance, Uploadable},
14    environment::{
15        AccessConfig, BankQuerier, BankSetter, ChainInfoOwned, ChainState, DefaultQueriers,
16        IndexResponse, StateInterface, TxHandler,
17    },
18    CwEnvError,
19};
20use cw_orch_daemon::{queriers::Node, read_network_config, DEFAULT_DEPLOYMENT, RUNTIME};
21use cw_utils::NativeBalance;
22use serde::Serialize;
23use tokio::runtime::Runtime;
24
25use crate::{contract::CloneTestingContract, queriers::bank::CloneBankQuerier};
26
27use super::state::MockState;
28
29pub type CloneTestingApp = App<BankKeeper, MockApiBech32>;
30
31/// Wrapper around a cw-multi-test [`App`](cw_multi_test::App) backend.
32///
33/// Stores a local state with a mapping of contract_id -> code_id/address
34///
35/// The state is customizable by implementing the [`StateInterface`] trait on a custom struct and providing it on the custom constructor.
36///
37/// ## Example
38/// ```
39/// # use cosmwasm_std::{Addr, coin, Uint128};
40/// use cw_orch_clone_testing::CloneTesting;
41/// use cw_orch_core::environment::TxHandler;
42///
43/// let chain = cw_orch_daemon::networks::JUNO_1;
44/// let mock: CloneTesting = CloneTesting::new(chain.clone()).unwrap();
45/// let sender = mock.sender();
46///
47/// // set a balance
48/// mock.set_balance(&sender, vec![coin(100u128, "token")]).unwrap();
49///
50/// // query the balance
51/// let balance: Uint128 = mock.query_balance(&sender, "token").unwrap();
52/// assert_eq!(balance.u128(), 100u128);
53/// ```
54///
55/// ## Example with custom state
56/// ```
57/// # use cosmwasm_std::{Addr, coin, Uint128};
58/// use cw_orch_clone_testing::CloneTesting;
59/// use cw_orch_core::environment::StateInterface;
60/// // We just use the MockState as an example here, but you can implement your own state struct.
61/// use cw_orch_clone_testing::MockState as CustomState;
62///
63/// let rt = tokio::runtime::Runtime::new().unwrap();
64/// let chain = cw_orch_daemon::networks::JUNO_1;
65/// let mock: CloneTesting = CloneTesting::new_custom(&rt, chain.clone(), CustomState::new(chain.clone().into(), "mock")).unwrap();
66/// ```
67#[derive(Clone)]
68pub struct CloneTesting<S: StateInterface = MockState> {
69    /// Chain data of the chain you want to fork
70    pub chain: ChainInfoOwned,
71    /// Address used for the operations.
72    pub sender: Addr,
73    /// Inner mutable state storage for contract addresses and code-ids
74    pub state: Rc<RefCell<S>>,
75    /// Inner mutable cw-multi-test app backend
76    pub app: Rc<RefCell<CloneTestingApp>>,
77}
78
79impl CloneTesting {
80    /// Ceates a new valid account
81    pub fn init_account(&self) -> Addr {
82        self.app.borrow_mut().next_address()
83    }
84
85    /// Set the bank balance of an address.
86    pub fn set_balance(
87        &self,
88        address: &Addr,
89        amount: Vec<cosmwasm_std::Coin>,
90    ) -> Result<(), CwEnvError> {
91        self.app
92            .borrow_mut()
93            .init_modules(|router, _, storage| router.bank.init_balance(storage, address, amount))
94            .map_err(Into::into)
95    }
96
97    /// Adds the bank balance of an address.
98    pub fn add_balance(
99        &self,
100        address: &Addr,
101        amount: Vec<cosmwasm_std::Coin>,
102    ) -> Result<(), CwEnvError> {
103        let b = self.query_all_balances(address)?;
104        let new_amount = NativeBalance(b) + NativeBalance(amount);
105        self.app
106            .borrow_mut()
107            .init_modules(|router, _, storage| {
108                router
109                    .bank
110                    .init_balance(storage, address, new_amount.into_vec())
111            })
112            .map_err(Into::into)
113    }
114
115    /// Set the balance for multiple coins at once.
116    pub fn set_balances(
117        &self,
118        balances: &[(&Addr, &[cosmwasm_std::Coin])],
119    ) -> Result<(), CwEnvError> {
120        self.app
121            .borrow_mut()
122            .init_modules(|router, _, storage| -> Result<(), CwEnvError> {
123                for (addr, coins) in balances {
124                    router.bank.init_balance(storage, addr, coins.to_vec())?;
125                }
126                Ok(())
127            })
128    }
129
130    /// Query the (bank) balance of a native token for and address.
131    /// Returns the amount of the native token.
132    pub fn query_balance(&self, address: &Addr, denom: &str) -> Result<Uint128, CwEnvError> {
133        Ok(self
134            .bank_querier()
135            .balance(address, Some(denom.to_string()))?[0]
136            .amount)
137    }
138
139    /// Fetch all the balances of an address.
140    pub fn query_all_balances(
141        &self,
142        address: &Addr,
143    ) -> Result<Vec<cosmwasm_std::Coin>, CwEnvError> {
144        self.bank_querier().balance(address, None)
145    }
146
147    pub fn upload_wasm<T: Uploadable + ContractInstance<CloneTesting>>(
148        &self,
149        contract: &T,
150    ) -> Result<<Self as TxHandler>::Response, CwEnvError> {
151        let mut file = std::fs::File::open(T::wasm(&self.chain).path())?;
152        let mut wasm = Vec::<u8>::new();
153        file.read_to_end(&mut wasm)?;
154        let code_id = self.app.borrow_mut().store_wasm_code(wasm);
155
156        contract.set_code_id(code_id);
157
158        // add contract code_id to events manually
159        let mut event = Event::new("store_code");
160        event = event.add_attribute("code_id", code_id.to_string());
161        let resp = AppResponse {
162            events: vec![event],
163            ..Default::default()
164        };
165        Ok(resp)
166    }
167}
168
169impl CloneTesting<MockState> {
170    /// Create a mock environment with the default mock state.
171    pub fn new(chain: impl Into<ChainInfoOwned>) -> Result<Self, CwEnvError> {
172        Self::new_with_runtime(&RUNTIME, chain)
173    }
174
175    /// Create a mock environment with the default mock state.
176    /// It uses a custom runtime object to control async requests
177    pub fn new_with_runtime(
178        rt: &Runtime,
179        chain: impl Into<ChainInfoOwned>,
180    ) -> Result<Self, CwEnvError> {
181        let chain_data = chain.into();
182        CloneTesting::new_custom(
183            rt,
184            chain_data.clone(),
185            MockState::new(chain_data, DEFAULT_DEPLOYMENT),
186        )
187    }
188
189    pub fn new_with_deployment_id(
190        rt: &Runtime,
191        chain: impl Into<ChainInfoOwned>,
192        deployment_id: &str,
193    ) -> Result<Self, CwEnvError> {
194        let chain_data = chain.into();
195        CloneTesting::new_custom(
196            rt,
197            chain_data.clone(),
198            MockState::new(chain_data, deployment_id),
199        )
200    }
201}
202
203impl<S: StateInterface> CloneTesting<S> {
204    /// Create a mock environment with a custom mock state.
205    /// The state is customizable by implementing the `StateInterface` trait on a custom struct and providing it on the custom constructor.
206    pub fn new_custom(
207        rt: &Runtime,
208        chain: impl Into<ChainInfoOwned>,
209        custom_state: S,
210    ) -> Result<Self, CwEnvError> {
211        let chain: ChainInfoOwned = chain.into();
212        let chain = if let Some(chain_info) = read_network_config(&chain.chain_id) {
213            chain.overwrite_with(chain_info)
214        } else {
215            chain
216        };
217        let state = Rc::new(RefCell::new(custom_state));
218
219        let pub_address_prefix = chain.network_info.pub_address_prefix.clone();
220        let remote_channel = RemoteChannel::new(
221            rt,
222            &chain
223                .grpc_urls
224                .iter()
225                .map(|s| s.as_str())
226                .collect::<Vec<_>>(),
227            &chain.chain_id,
228            &chain.network_info.pub_address_prefix,
229        )
230        .unwrap();
231
232        let wasm = WasmKeeper::<Empty, Empty>::new()
233            .with_remote(remote_channel.clone())
234            .with_address_generator(MockAddressGenerator);
235
236        let bank = BankKeeper::new().with_remote(remote_channel.clone());
237
238        // We update the block_height
239        let block_info = remote_channel
240            .rt
241            .block_on(Node::new_async(remote_channel.channel.clone())._block_info())
242            .unwrap();
243
244        // Finally we instantiate a new app
245        let app = AppBuilder::default()
246            .with_wasm(wasm)
247            .with_bank(bank)
248            .with_api(MockApiBech32::new(&pub_address_prefix))
249            .with_block(block_info)
250            .with_remote(remote_channel.clone());
251
252        let app = Rc::new(RefCell::new(app.build(|_, _, _| {})?));
253        let sender = app.borrow_mut().next_address();
254
255        Ok(Self {
256            chain,
257            sender: sender.clone(),
258            state,
259            app,
260        })
261    }
262
263    pub fn storage_analysis(&self) -> StorageAnalyzer {
264        StorageAnalyzer::new(&self.app.borrow()).unwrap()
265    }
266}
267
268impl<S: StateInterface> ChainState for CloneTesting<S> {
269    type Out = Rc<RefCell<S>>;
270
271    fn state(&self) -> Self::Out {
272        self.state.clone()
273    }
274
275    fn can_load_state_from_state_file(&self) -> bool {
276        true
277    }
278}
279
280// Execute on the test chain, returns test response type
281impl<S: StateInterface> TxHandler for CloneTesting<S> {
282    type Response = AppResponse;
283    type Error = CwEnvError;
284    type ContractSource = Box<dyn Contract<Empty, Empty>>;
285    type Sender = Addr;
286
287    fn sender(&self) -> &cosmwasm_std::Addr {
288        &self.sender
289    }
290
291    fn sender_addr(&self) -> Addr {
292        self.sender.clone()
293    }
294
295    fn set_sender(&mut self, sender: Self::Sender) {
296        self.sender = sender;
297    }
298
299    fn upload<T: Uploadable>(&self, _contract: &T) -> Result<Self::Response, CwEnvError> {
300        let wrapper_contract = CloneTestingContract::new(T::wrapper());
301        let code_id = self
302            .app
303            .borrow_mut()
304            .store_code_with_creator(self.sender_addr(), Box::new(wrapper_contract));
305        // add contract code_id to events manually
306        let mut event = Event::new("store_code");
307        event = event.add_attribute("code_id", code_id.to_string());
308        let resp = AppResponse {
309            events: vec![event],
310            ..Default::default()
311        };
312        Ok(resp)
313    }
314
315    fn upload_with_access_config<T: Uploadable>(
316        &self,
317        contract_source: &T,
318        _access_config: Option<AccessConfig>,
319    ) -> Result<Self::Response, Self::Error> {
320        log::debug!("Uploading with access is not enforced when using Clone Testing");
321        self.upload(contract_source)
322    }
323
324    fn execute<E: Serialize + Debug>(
325        &self,
326        exec_msg: &E,
327        coins: &[cosmwasm_std::Coin],
328        contract_address: &Addr,
329    ) -> Result<Self::Response, CwEnvError> {
330        self.app
331            .borrow_mut()
332            .execute_contract(
333                self.sender.clone(),
334                contract_address.to_owned(),
335                exec_msg,
336                coins,
337            )
338            .map_err(From::from)
339            .map(Into::into)
340    }
341
342    fn instantiate<I: Serialize + Debug>(
343        &self,
344        code_id: u64,
345        init_msg: &I,
346        label: Option<&str>,
347        admin: Option<&Addr>,
348        coins: &[cosmwasm_std::Coin],
349    ) -> Result<Self::Response, CwEnvError> {
350        let addr = self.app.borrow_mut().instantiate_contract(
351            code_id,
352            self.sender.clone(),
353            init_msg,
354            coins,
355            label.unwrap_or("contract_init"),
356            admin.map(|a| a.to_string()),
357        )?;
358        // add contract address to events manually
359        let mut event = Event::new("instantiate");
360        event = event.add_attribute("_contract_address", addr);
361        let resp = AppResponse {
362            events: vec![event],
363            ..Default::default()
364        };
365        Ok(resp)
366    }
367
368    fn migrate<M: Serialize + Debug>(
369        &self,
370        migrate_msg: &M,
371        new_code_id: u64,
372        contract_address: &Addr,
373    ) -> Result<Self::Response, CwEnvError> {
374        self.app
375            .borrow_mut()
376            .migrate_contract(
377                self.sender.clone(),
378                contract_address.clone(),
379                migrate_msg,
380                new_code_id,
381            )
382            .map_err(From::from)
383            .map(Into::into)
384    }
385
386    fn instantiate2<I: Serialize + Debug>(
387        &self,
388        code_id: u64,
389        init_msg: &I,
390        label: Option<&str>,
391        admin: Option<&Addr>,
392        coins: &[cosmwasm_std::Coin],
393        salt: Binary,
394    ) -> Result<Self::Response, Self::Error> {
395        let resp = self.app.borrow_mut().execute(
396            self.sender.clone(),
397            CosmosMsg::Wasm(WasmMsg::Instantiate2 {
398                admin: admin.map(|a| a.to_string()),
399                code_id,
400                label: label.unwrap_or("contract_init").to_string(),
401                msg: to_json_binary(init_msg)?,
402                funds: coins.to_vec(),
403                salt,
404            }),
405        )?;
406
407        let app_resp = AppResponse {
408            events: resp.events,
409            data: resp.data,
410        };
411
412        Ok(app_resp)
413    }
414
415    fn bank_send(
416        &self,
417        receiver: &Addr,
418        amount: &[cosmwasm_std::Coin],
419    ) -> Result<Self::Response, Self::Error> {
420        self.app
421            .borrow_mut()
422            .execute(
423                self.sender.clone(),
424                BankMsg::Send {
425                    to_address: receiver.to_string(),
426                    amount: amount.to_vec(),
427                }
428                .into(),
429            )
430            .map_err(From::from)
431            .map(Into::into)
432    }
433}
434
435/// Custom AppResponse type for working with the IndexResponse trait
436#[derive(Default, Clone, Debug)]
437pub struct AppResponse {
438    pub events: Vec<Event>,
439    pub data: Option<Binary>,
440}
441
442impl From<clone_cw_multi_test::AppResponse> for AppResponse {
443    fn from(value: clone_cw_multi_test::AppResponse) -> Self {
444        AppResponse {
445            events: value.events,
446            data: value.data,
447        }
448    }
449}
450
451impl From<AppResponse> for clone_cw_multi_test::AppResponse {
452    fn from(value: AppResponse) -> Self {
453        clone_cw_multi_test::AppResponse {
454            events: value.events,
455            data: value.data,
456        }
457    }
458}
459
460impl IndexResponse for AppResponse {
461    fn events(&self) -> Vec<Event> {
462        self.events.clone()
463    }
464
465    fn data(&self) -> Option<Binary> {
466        self.data.clone()
467    }
468
469    fn event_attr_value(&self, event_type: &str, attr_key: &str) -> StdResult<String> {
470        for event in &self.events {
471            if event.ty == event_type {
472                for attr in &event.attributes {
473                    if attr.key == attr_key {
474                        return Ok(attr.value.clone());
475                    }
476                }
477            }
478        }
479        Err(StdError::generic_err(format!(
480            "missing combination (event: {}, attribute: {})",
481            event_type, attr_key
482        )))
483    }
484
485    fn event_attr_values(&self, event_type: &str, attr_key: &str) -> Vec<String> {
486        let mut all_results = vec![];
487
488        for event in &self.events {
489            if event.ty == event_type {
490                for attr in &event.attributes {
491                    if attr.key == attr_key {
492                        all_results.push(attr.value.clone());
493                    }
494                }
495            }
496        }
497        all_results
498    }
499}
500
501impl BankSetter for CloneTesting {
502    type T = CloneBankQuerier;
503    fn set_balance(
504        &mut self,
505        address: &Addr,
506        amount: Vec<Coin>,
507    ) -> Result<(), <Self as TxHandler>::Error> {
508        (*self).set_balance(&Addr::unchecked(address), amount)
509    }
510}
511
512#[cfg(test)]
513mod test {
514    use crate::core::*;
515    use clone_cw_multi_test::LOCAL_RUST_CODE_OFFSET;
516    use cosmwasm_std::{
517        to_json_binary, Addr, Coin, Deps, DepsMut, Env, MessageInfo, Response, Uint128,
518    };
519    use cw20::{BalanceResponse, MinterResponse};
520    use cw_orch_core::contract::WasmPath;
521    use cw_orch_core::environment::QueryHandler;
522    use cw_orch_daemon::networks::JUNO_1;
523    use cw_orch_mock::cw_multi_test::{Contract as MockContract, ContractWrapper};
524    use speculoos::prelude::*;
525
526    pub struct MockCw20;
527
528    fn execute(
529        _deps: DepsMut,
530        _env: Env,
531        _info: MessageInfo,
532        msg: cw20::Cw20ExecuteMsg,
533    ) -> Result<Response, cw20_base::ContractError> {
534        match msg {
535            cw20::Cw20ExecuteMsg::Mint { recipient, amount } => Ok(Response::default()
536                .add_attribute("action", "mint")
537                .add_attribute("recipient", recipient)
538                .add_attribute("amount", amount)),
539            _ => unimplemented!(),
540        }
541    }
542
543    fn query(_deps: Deps, _env: Env, msg: cw20_base::msg::QueryMsg) -> StdResult<Binary> {
544        match msg {
545            cw20_base::msg::QueryMsg::Balance { address: _ } => {
546                Ok(to_json_binary::<BalanceResponse>(&BalanceResponse {
547                    balance: Uint128::from(100u128),
548                })
549                .unwrap())
550            }
551            _ => unimplemented!(),
552        }
553    }
554    impl Uploadable for MockCw20 {
555        fn wasm(_chain: &ChainInfoOwned) -> WasmPath {
556            unimplemented!()
557        }
558
559        fn wrapper() -> Box<dyn MockContract<Empty, Empty>> {
560            Box::new(
561                ContractWrapper::new(execute, cw20_base::contract::instantiate, query)
562                    .with_migrate(cw20_base::contract::migrate),
563            )
564        }
565    }
566
567    #[test]
568    fn mock() -> anyhow::Result<()> {
569        let amount = 1000000u128;
570        let denom = "uosmo";
571        let chain_info = JUNO_1;
572
573        let chain = CloneTesting::new(chain_info)?;
574
575        let sender = chain.sender_addr();
576        let recipient = &chain.init_account();
577
578        chain
579            .set_balance(recipient, vec![Coin::new(amount, denom)])
580            .unwrap();
581        let balance = chain.query_balance(recipient, denom).unwrap();
582
583        asserting("address balance amount is correct")
584            .that(&amount)
585            .is_equal_to(balance.u128());
586
587        asserting("sender is correct")
588            .that(&sender)
589            .is_equal_to(chain.sender_addr());
590
591        let init_res = chain.upload(&MockCw20).unwrap();
592        let code_id = (1 + LOCAL_RUST_CODE_OFFSET) as u64;
593        asserting("contract initialized properly")
594            .that(&init_res.events[0].attributes[0].value)
595            .is_equal_to(code_id.to_string());
596
597        let init_msg = cw20_base::msg::InstantiateMsg {
598            name: String::from("Token"),
599            symbol: String::from("TOK"),
600            decimals: 6u8,
601            initial_balances: vec![],
602            mint: Some(MinterResponse {
603                minter: sender.to_string(),
604                cap: None,
605            }),
606            marketing: None,
607        };
608        let init_res = chain
609            .instantiate(code_id, &init_msg, None, Some(&sender), &[])
610            .unwrap();
611
612        let contract_address = Addr::unchecked(&init_res.events[0].attributes[0].value);
613
614        let exec_res = chain
615            .execute(
616                &cw20_base::msg::ExecuteMsg::Mint {
617                    recipient: recipient.to_string(),
618                    amount: Uint128::from(100u128),
619                },
620                &[],
621                &contract_address,
622            )
623            .unwrap();
624
625        asserting("that exect passed on correctly")
626            .that(&exec_res.events[1].attributes[1].value)
627            .is_equal_to(String::from("mint"));
628
629        let query_res = chain
630            .query::<cw20_base::msg::QueryMsg, BalanceResponse>(
631                &cw20_base::msg::QueryMsg::Balance {
632                    address: recipient.to_string(),
633                },
634                &contract_address,
635            )
636            .unwrap();
637
638        asserting("that query passed on correctly")
639            .that(&query_res.balance)
640            .is_equal_to(Uint128::from(100u128));
641
642        let migration_res = chain.migrate(&Empty {}, code_id, &contract_address);
643        asserting("that migration passed on correctly")
644            .that(&migration_res)
645            .is_ok();
646
647        Ok(())
648    }
649
650    #[test]
651    fn custom_mock_env() -> anyhow::Result<()> {
652        let amount = 1000000u128;
653        let denom = "uosmo";
654        let chain = JUNO_1;
655
656        let rt = Runtime::new().unwrap();
657        let mock_state = MockState::new(JUNO_1.into(), "default_id");
658
659        let chain: CloneTesting = CloneTesting::<_>::new_custom(&rt, chain, mock_state)?;
660        let recipient = chain.init_account();
661
662        chain
663            .set_balances(&[(&recipient, &[Coin::new(amount, denom)])])
664            .unwrap();
665
666        let balances = chain.query_all_balances(&recipient).unwrap();
667        asserting("recipient balances length is 1")
668            .that(&balances.len())
669            .is_equal_to(1);
670
671        Ok(())
672    }
673
674    #[test]
675    fn state_interface() {
676        let contract_id = "my_contract";
677        let code_id = 1u64;
678        let address = &Addr::unchecked("TEST_ADDR");
679        let mut mock_state = Rc::new(RefCell::new(MockState::new(JUNO_1.into(), "default_id")));
680
681        mock_state.set_address(contract_id, address);
682        asserting!("that address has been set")
683            .that(&address)
684            .is_equal_to(&mock_state.get_address(contract_id).unwrap());
685
686        mock_state.set_code_id(contract_id, code_id);
687        asserting!("that code_id has been set")
688            .that(&code_id)
689            .is_equal_to(mock_state.get_code_id(contract_id).unwrap());
690
691        asserting!("that total code_ids is 1")
692            .that(&mock_state.get_all_code_ids().unwrap().len())
693            .is_greater_than_or_equal_to(1);
694
695        asserting!("that total addresses is 1")
696            .that(&mock_state.get_all_addresses().unwrap().len())
697            .is_greater_than_or_equal_to(1);
698    }
699
700    #[test]
701    fn add_balance() -> anyhow::Result<()> {
702        let amount = 1000000u128;
703        let denom_1 = "uosmo";
704        let denom_2 = "osmou";
705        let chain_info = JUNO_1;
706
707        let chain = CloneTesting::new(chain_info)?;
708        let recipient = &chain.init_account();
709
710        chain
711            .add_balance(recipient, vec![Coin::new(amount, denom_1)])
712            .unwrap();
713        chain
714            .add_balance(recipient, vec![Coin::new(amount, denom_2)])
715            .unwrap();
716
717        let balances = chain.query_all_balances(recipient).unwrap();
718        asserting("recipient balances added")
719            .that(&balances)
720            .contains_all_of(&[&Coin::new(amount, denom_1), &Coin::new(amount, denom_2)]);
721        Ok(())
722    }
723}