Skip to main content

cw_orch_daemon/
live_mock.rs

1//! Live mock is a mock that uses a live chain to query for data.
2//! It can be used to do chain-backed unit-testing. It can't be used for state-changing operations.
3
4use crate::queriers::Bank;
5use crate::queriers::CosmWasm;
6use crate::queriers::Staking;
7use crate::RUNTIME;
8use cosmwasm_std::testing::{MockApi, MockStorage};
9use cosmwasm_std::Addr;
10use cosmwasm_std::AllBalanceResponse;
11use cosmwasm_std::BalanceResponse;
12use cosmwasm_std::BankQuery;
13use cosmwasm_std::Binary;
14use cosmwasm_std::Delegation;
15use cosmwasm_std::Empty;
16use cosmwasm_std::StakingQuery;
17use cosmwasm_std::{
18    from_json, to_json_binary, Coin, ContractResult, OwnedDeps, Querier, QuerierResult,
19    QueryRequest, SystemError, SystemResult, Uint128, WasmQuery,
20};
21use cosmwasm_std::{AllDelegationsResponse, BondedDenomResponse};
22use cw_orch_core::environment::BankQuerier;
23use cw_orch_core::environment::ChainInfoOwned;
24use cw_orch_core::environment::WasmQuerier;
25use std::marker::PhantomData;
26use std::str::FromStr;
27use tonic::transport::Channel;
28
29use crate::channel::GrpcChannel;
30
31fn to_cosmwasm_coin(c: cosmrs::proto::cosmos::base::v1beta1::Coin) -> Coin {
32    Coin {
33        amount: Uint128::from_str(&c.amount).unwrap(),
34        denom: c.denom,
35    }
36}
37
38const QUERIER_ERROR: &str =
39    "Only Bank balances and Wasm (raw + smart) and Some staking queries are covered for now";
40
41/// mock_dependencies is a drop-in replacement for cosmwasm_std::testing::mock_dependencies
42/// this uses our CustomQuerier.
43pub fn mock_dependencies(
44    chain_info: ChainInfoOwned,
45) -> OwnedDeps<MockStorage, MockApi, WasmMockQuerier> {
46    let custom_querier: WasmMockQuerier = WasmMockQuerier::new(chain_info);
47
48    OwnedDeps {
49        storage: MockStorage::default(),
50        api: MockApi::default(),
51        querier: custom_querier,
52        custom_query_type: PhantomData,
53    }
54}
55
56/// Querier struct that fetches queries on-chain directly
57pub struct WasmMockQuerier {
58    channel: Channel,
59}
60
61impl Querier for WasmMockQuerier {
62    fn raw_query(&self, bin_request: &[u8]) -> QuerierResult {
63        // MockQuerier doesn't support Custom, so we ignore it completely here
64        let request: QueryRequest<Empty> = match from_json(bin_request) {
65            Ok(v) => v,
66            Err(e) => {
67                return SystemResult::Err(SystemError::InvalidRequest {
68                    error: format!("Parsing query request: {}", e),
69                    request: bin_request.into(),
70                })
71            }
72        };
73        self.handle_query(&request)
74    }
75}
76
77impl WasmMockQuerier {
78    /// Function used to handle a query and customize the query behavior
79    /// This implements some queries by querying an actual node for the responses
80    pub fn handle_query(&self, request: &QueryRequest<Empty>) -> QuerierResult {
81        let handle = RUNTIME.handle();
82        match &request {
83            QueryRequest::Wasm(x) => {
84                let querier = CosmWasm::new_sync(self.channel.clone(), handle);
85                match x {
86                    WasmQuery::Smart { contract_addr, msg } => {
87                        // We forward the request to the cosmwasm querier
88
89                        let query_result: Result<Binary, _> = handle
90                            .block_on(
91                                querier
92                                    ._contract_state(&Addr::unchecked(contract_addr), msg.to_vec()),
93                            )
94                            .map(|query_result| query_result.into());
95                        SystemResult::Ok(ContractResult::from(query_result))
96                    }
97                    WasmQuery::Raw { contract_addr, key } => {
98                        // We forward the request to the cosmwasm querie
99                        let query_result = querier
100                            .raw_query(&Addr::unchecked(contract_addr), key.to_vec())
101                            .map(|query_result| query_result.into());
102
103                        SystemResult::Ok(ContractResult::from(query_result))
104                    }
105                    _ => SystemResult::Err(SystemError::InvalidRequest {
106                        error: QUERIER_ERROR.to_string(),
107                        request: to_json_binary(&request).unwrap(),
108                    }),
109                }
110            }
111            QueryRequest::Bank(x) => {
112                let querier = Bank {
113                    channel: self.channel.clone(),
114                    rt_handle: Some(handle.clone()),
115                };
116                match x {
117                    BankQuery::Balance { address, denom } => {
118                        let query_result = querier
119                            .balance(&Addr::unchecked(address), Some(denom.clone()))
120                            .map(|result| {
121                                to_json_binary(&BalanceResponse::new(result[0].clone())).unwrap()
122                            });
123                        SystemResult::Ok(ContractResult::from(query_result))
124                    }
125                    BankQuery::AllBalances { address } => {
126                        let query_result = querier
127                            .balance(&Addr::unchecked(address), None)
128                            .map(AllBalanceResponse::new)
129                            .map(|query_result| to_json_binary(&query_result))
130                            .unwrap();
131                        SystemResult::Ok(ContractResult::from(query_result))
132                    }
133                    _ => SystemResult::Err(SystemError::InvalidRequest {
134                        error: QUERIER_ERROR.to_string(),
135                        request: to_json_binary(&request).unwrap(),
136                    }),
137                }
138            }
139            QueryRequest::Staking(x) => {
140                let querier = Staking::new_async(self.channel.clone());
141                match x {
142                    StakingQuery::BondedDenom {} => {
143                        let query_result = handle
144                            .block_on(querier._params())
145                            .map(|result| {
146                                BondedDenomResponse::new(result.params.unwrap().bond_denom)
147                            })
148                            .map(|query_result| to_json_binary(&query_result))
149                            .unwrap();
150                        SystemResult::Ok(ContractResult::from(query_result))
151                    }
152                    // This query is not perfect. I guess that on_chain you should be able to get ALL delegations and not a paginated result
153                    // TODO, do better here
154                    StakingQuery::AllDelegations { delegator } => {
155                        let query_result = handle
156                            .block_on(
157                                querier._delegator_delegations(&Addr::unchecked(delegator), None),
158                            )
159                            .map(|result| {
160                                AllDelegationsResponse::new(
161                                    result
162                                        .delegation_responses
163                                        .into_iter()
164                                        .filter_map(|delegation| {
165                                            delegation.delegation.map(|d| {
166                                                Delegation::new(
167                                                    Addr::unchecked(d.delegator_address),
168                                                    d.validator_address,
169                                                    to_cosmwasm_coin(delegation.balance.unwrap()),
170                                                )
171                                            })
172                                        })
173                                        .collect(),
174                                )
175                            })
176                            .map(|query_result| to_json_binary(&query_result))
177                            .unwrap();
178                        SystemResult::Ok(ContractResult::from(query_result))
179                    }
180                    _ => todo!(),
181                }
182            }
183            _ => SystemResult::Err(SystemError::InvalidRequest {
184                error: QUERIER_ERROR.to_string(),
185                request: to_json_binary(&request).unwrap(),
186            }),
187        }
188    }
189}
190
191impl WasmMockQuerier {
192    /// Creates a querier from chain information
193    pub fn new(chain: ChainInfoOwned) -> Self {
194        let channel = RUNTIME
195            .block_on(GrpcChannel::connect(
196                &chain.grpc_urls,
197                chain.chain_id.as_str(),
198            ))
199            .unwrap();
200
201        WasmMockQuerier { channel }
202    }
203}
204
205#[cfg(test)]
206mod tests {
207
208    use super::*;
209    use crate::networks::JUNO_1;
210
211    use super::mock_dependencies;
212
213    #[test]
214    fn bank_balance_querier() -> Result<(), anyhow::Error> {
215        let address = "juno1rkhrfuq7k2k68k0hctrmv8efyxul6tgn8hny6y";
216
217        let deps = mock_dependencies(JUNO_1.into());
218        let deps_ref = deps.as_ref();
219        let _response: BalanceResponse =
220            deps_ref
221                .querier
222                .query(&QueryRequest::Bank(BankQuery::Balance {
223                    address: address.to_string(),
224                    denom: "ujuno".to_string(),
225                }))?;
226        // We can't really test that response, but it has to unwrap at least !
227
228        Ok(())
229    }
230
231    #[test]
232    fn bank_all_balances_querier() -> Result<(), anyhow::Error> {
233        let address = "juno1rkhrfuq7k2k68k0hctrmv8efyxul6tgn8hny6y";
234
235        let deps = mock_dependencies(JUNO_1.into());
236        let deps_ref = deps.as_ref();
237        let _response: AllBalanceResponse =
238            deps_ref
239                .querier
240                .query(&QueryRequest::Bank(BankQuery::AllBalances {
241                    address: address.to_string(),
242                }))?;
243        // We can't really test that response, but it has to unwrap at least !
244        Ok(())
245    }
246}