clone_cw_multi_test/wasm_emulation/query/
mock_querier.rs

1use std::collections::HashMap;
2
3use crate::wasm_emulation::channel::RemoteChannel;
4use crate::wasm_emulation::query::bank::BankQuerier;
5use crate::wasm_emulation::query::staking::StakingQuerier;
6use crate::wasm_emulation::query::wasm::WasmQuerier;
7
8use cosmwasm_std::CustomMsg;
9use cosmwasm_std::Env;
10use cosmwasm_vm::BackendResult;
11use cosmwasm_vm::GasInfo;
12
13use serde::de::DeserializeOwned;
14
15use cosmwasm_std::Binary;
16use cosmwasm_std::Coin;
17
18use cosmwasm_std::SystemError;
19
20use cosmwasm_std::from_json;
21use cosmwasm_std::{ContractResult, SystemResult};
22use cosmwasm_std::{CustomQuery, QueryRequest};
23use cosmwasm_std::{FullDelegation, Validator};
24
25use cosmwasm_std::Attribute;
26use cosmwasm_std::QuerierResult;
27
28use crate::wasm_emulation::input::QuerierStorage;
29use crate::Contract;
30
31use super::gas::GAS_COST_QUERY_ERROR;
32
33#[derive(Clone)]
34pub struct LocalForkedState<ExecC, QueryC> {
35    pub contracts: HashMap<usize, *mut dyn Contract<ExecC, QueryC>>,
36    pub env: Env,
37}
38
39#[derive(Clone)]
40pub struct ForkState<ExecC, QueryC>
41where
42    QueryC: CustomQuery + DeserializeOwned + 'static,
43    ExecC: CustomMsg + 'static,
44{
45    pub remote: RemoteChannel,
46    /// Only query function right now, but we might pass along the whole application state to avoid stargate queries
47    pub local_state: LocalForkedState<ExecC, QueryC>,
48    pub querier_storage: QuerierStorage,
49}
50
51pub type QueryResultWithGas = (QuerierResult, GasInfo);
52
53/// The same type as cosmwasm-std's QuerierResult, but easier to reuse in
54/// cosmwasm-vm. It might diverge from QuerierResult at some point.
55pub type MockQuerierCustomHandlerResult = SystemResult<ContractResult<Binary>>;
56
57/// MockQuerier holds an immutable table of bank balances
58/// and configurable handlers for Wasm queries and custom queries.
59pub struct MockQuerier<
60    ExecC: CustomMsg + DeserializeOwned + 'static,
61    QueryC: CustomQuery + DeserializeOwned + 'static,
62> {
63    bank: BankQuerier,
64
65    staking: StakingQuerier,
66    wasm: WasmQuerier<ExecC, QueryC>,
67
68    //Box<dyn Fn(Deps<'_, C>, Env, Vec<u8>) -> Result<Binary, anyhow::Error>>, //fn(deps: Deps<C>, env: Env, msg: Vec<u8>) -> Result<Binary, anyhow::Error>,
69    /// A handler to handle custom queries. This is set to a dummy handler that
70    /// always errors by default. Update it via `with_custom_handler`.
71    ///
72    /// Use box to avoid the need of another generic type
73    custom_handler: Box<dyn for<'a> Fn(&'a QueryC) -> QueryResultWithGas>,
74    remote: RemoteChannel,
75}
76
77impl<
78        ExecC: CustomMsg + DeserializeOwned + 'static,
79        QueryC: CustomQuery + DeserializeOwned + 'static,
80    > MockQuerier<ExecC, QueryC>
81{
82    pub fn new(fork_state: ForkState<ExecC, QueryC>) -> Self {
83        // We create query_closures for all local_codes
84
85        MockQuerier {
86            bank: BankQuerier::new(
87                fork_state.remote.clone(),
88                fork_state.querier_storage.bank.storage.clone(),
89            ),
90
91            staking: StakingQuerier::default(),
92            wasm: WasmQuerier::new(fork_state.clone()),
93            // strange argument notation suggested as a workaround here: https://github.com/rust-lang/rust/issues/41078#issuecomment-294296365
94            custom_handler: Box::from(|_: &_| -> QueryResultWithGas {
95                (
96                    SystemResult::Err(SystemError::UnsupportedRequest {
97                        kind: "custom".to_string(),
98                    }),
99                    GasInfo::free(),
100                )
101            }),
102            remote: fork_state.remote,
103        }
104    }
105
106    // set a new balance for the given address and return the old balance
107    pub fn update_balance(
108        &mut self,
109        addr: impl Into<String>,
110        balance: Vec<Coin>,
111    ) -> Option<Vec<Coin>> {
112        self.bank.update_balance(addr, balance)
113    }
114
115    pub fn update_staking(
116        &mut self,
117        denom: &str,
118        validators: &[Validator],
119        delegations: &[FullDelegation],
120    ) {
121        self.staking = StakingQuerier::new(denom, validators, delegations);
122    }
123
124    pub fn with_custom_handler<CH>(mut self, handler: CH) -> Self
125    where
126        CH: Fn(&QueryC) -> QueryResultWithGas + 'static,
127    {
128        self.custom_handler = Box::from(handler);
129        self
130    }
131}
132
133impl<
134        ExecC: CustomMsg + DeserializeOwned + 'static,
135        QueryC: CustomQuery + DeserializeOwned + 'static,
136    > cosmwasm_vm::Querier for MockQuerier<ExecC, QueryC>
137{
138    fn query_raw(
139        &self,
140        bin_request: &[u8],
141        _gas_limit: u64,
142    ) -> BackendResult<SystemResult<ContractResult<Binary>>> {
143        let request: QueryRequest<QueryC> = match from_json(bin_request) {
144            Ok(v) => v,
145            Err(e) => {
146                return (
147                    Ok(SystemResult::Err(SystemError::InvalidRequest {
148                        error: format!("Parsing query request: {}", e),
149                        request: bin_request.into(),
150                    })),
151                    GasInfo::with_externally_used(GAS_COST_QUERY_ERROR),
152                )
153            }
154        };
155        let result = self.handle_query(&request);
156
157        (Ok(result.0), result.1)
158    }
159}
160
161impl<
162        ExecC: CustomMsg + DeserializeOwned + 'static,
163        QueryC: CustomQuery + DeserializeOwned + 'static,
164    > cosmwasm_std::Querier for MockQuerier<ExecC, QueryC>
165{
166    fn raw_query(&self, bin_request: &[u8]) -> SystemResult<ContractResult<Binary>> {
167        let request: QueryRequest<QueryC> = match from_json(bin_request) {
168            Ok(v) => v,
169            Err(e) => {
170                return SystemResult::Err(SystemError::InvalidRequest {
171                    error: format!("Parsing query request: {}", e),
172                    request: bin_request.into(),
173                })
174            }
175        };
176        let result = self.handle_query(&request);
177
178        result.0
179    }
180}
181
182impl<
183        ExecC: CustomMsg + DeserializeOwned + 'static,
184        QueryC: CustomQuery + DeserializeOwned + 'static,
185    > MockQuerier<ExecC, QueryC>
186{
187    pub fn handle_query(&self, request: &QueryRequest<QueryC>) -> QueryResultWithGas {
188        match &request {
189            QueryRequest::Bank(bank_query) => self.bank.query(bank_query),
190            QueryRequest::Custom(custom_query) => (*self.custom_handler)(custom_query),
191
192            QueryRequest::Staking(staking_query) => self.staking.query(staking_query),
193            QueryRequest::Wasm(msg) => self.wasm.query(self.remote.clone(), msg),
194            #[allow(deprecated)]
195            QueryRequest::Stargate { .. } => (
196                SystemResult::Err(SystemError::UnsupportedRequest {
197                    kind: "Stargate".to_string(),
198                }),
199                GasInfo::with_externally_used(GAS_COST_QUERY_ERROR),
200            ),
201            QueryRequest::Grpc(_req) => (
202                SystemResult::Err(SystemError::UnsupportedRequest {
203                    kind: "Stargate".to_string(),
204                }),
205                GasInfo::with_externally_used(GAS_COST_QUERY_ERROR),
206            ),
207            &_ => panic!("Query Type Not implemented"),
208        }
209    }
210}
211
212pub fn digit_sum(input: &[u8]) -> usize {
213    input.iter().fold(0, |sum, val| sum + (*val as usize))
214}
215
216/// Only for test code. This bypasses assertions in new, allowing us to create _*
217/// Attributes to simulate responses from the blockchain
218pub fn mock_wasmd_attr(key: impl Into<String>, value: impl Into<String>) -> Attribute {
219    Attribute {
220        key: key.into(),
221        value: value.into(),
222    }
223}
clone_cw_multi_test/wasm_emulation/query/mock_querier.rs

clone_cw_multi_test/wasm_emulation/query/
mock_querier.rs
