cosmwasm_vm/testing/
querier.rs

1use serde::de::DeserializeOwned;
2
3use cosmwasm_std::testing::{MockQuerier as StdMockQuerier, MockQuerierCustomHandlerResult};
4use cosmwasm_std::{
5    to_json_binary, to_json_vec, Binary, Coin, ContractResult, CustomQuery, Empty, Querier as _,
6    QueryRequest, SystemError, SystemResult,
7};
8
9use crate::{BackendError, BackendResult, GasInfo, Querier};
10
11const GAS_COST_QUERY_FLAT: u64 = 100_000;
12/// Gas per request byte
13const GAS_COST_QUERY_REQUEST_MULTIPLIER: u64 = 0;
14/// Gas per response byte
15const GAS_COST_QUERY_RESPONSE_MULTIPLIER: u64 = 100;
16
17/// MockQuerier holds an immutable table of bank balances
18pub struct MockQuerier<C: CustomQuery + DeserializeOwned = Empty> {
19    querier: StdMockQuerier<C>,
20}
21
22impl<C: CustomQuery + DeserializeOwned> MockQuerier<C> {
23    pub fn new(balances: &[(&str, &[Coin])]) -> Self {
24        MockQuerier {
25            querier: StdMockQuerier::new(balances),
26        }
27    }
28
29    /// Set a new balance for the given address and return the old balance
30    pub fn update_balance(
31        &mut self,
32        addr: impl Into<String>,
33        balance: Vec<Coin>,
34    ) -> Option<Vec<Coin>> {
35        self.querier.bank.update_balance(addr, balance)
36    }
37
38    #[cfg(feature = "staking")]
39    pub fn update_staking(
40        &mut self,
41        denom: &str,
42        validators: &[cosmwasm_std::Validator],
43        delegations: &[cosmwasm_std::FullDelegation],
44    ) {
45        self.querier.staking.update(denom, validators, delegations);
46    }
47
48    pub fn update_wasm<WH>(&mut self, handler: WH)
49    where
50        WH: Fn(&cosmwasm_std::WasmQuery) -> cosmwasm_std::QuerierResult + 'static,
51    {
52        self.querier.update_wasm(handler)
53    }
54
55    pub fn with_custom_handler<CH>(mut self, handler: CH) -> Self
56    where
57        CH: Fn(&C) -> MockQuerierCustomHandlerResult + 'static,
58    {
59        self.querier = self.querier.with_custom_handler(handler);
60        self
61    }
62}
63
64impl<C: CustomQuery + DeserializeOwned> Querier for MockQuerier<C> {
65    fn query_raw(
66        &self,
67        bin_request: &[u8],
68        gas_limit: u64,
69    ) -> BackendResult<SystemResult<ContractResult<Binary>>> {
70        let response = self.querier.raw_query(bin_request);
71        let gas_info = GasInfo::with_externally_used(
72            GAS_COST_QUERY_FLAT
73                + (GAS_COST_QUERY_REQUEST_MULTIPLIER * (bin_request.len() as u64))
74                + (GAS_COST_QUERY_RESPONSE_MULTIPLIER
75                    * (to_json_binary(&response).unwrap().len() as u64)),
76        );
77
78        // In a production implementation, this should stop the query execution in the middle of the computation.
79        // Thus no query response is returned to the caller.
80        if gas_info.externally_used > gas_limit {
81            return (Err(BackendError::out_of_gas()), gas_info);
82        }
83
84        // We don't use FFI in the mock implementation, so BackendResult is always Ok() regardless of error on other levels
85        (Ok(response), gas_info)
86    }
87}
88
89impl MockQuerier {
90    pub fn query<C: CustomQuery>(
91        &self,
92        request: &QueryRequest<C>,
93        gas_limit: u64,
94    ) -> BackendResult<SystemResult<ContractResult<Binary>>> {
95        // encode the request, then call raw_query
96        let request_binary = match to_json_vec(request) {
97            Ok(raw) => raw,
98            Err(err) => {
99                let gas_info = GasInfo::with_externally_used(err.to_string().len() as u64);
100                return (
101                    Ok(SystemResult::Err(SystemError::InvalidRequest {
102                        error: format!("Serializing query request: {err}"),
103                        request: b"N/A".into(),
104                    })),
105                    gas_info,
106                );
107            }
108        };
109        self.query_raw(&request_binary, gas_limit)
110    }
111}
112
113#[cfg(test)]
114mod tests {
115    use super::*;
116    use cosmwasm_std::{coin, from_json, BalanceResponse, BankQuery};
117
118    const DEFAULT_QUERY_GAS_LIMIT: u64 = 300_000;
119
120    #[test]
121    fn query_raw_fails_when_out_of_gas() {
122        let addr = String::from("foobar");
123        let balance = vec![coin(123, "ELF"), coin(777, "FLY")];
124        let querier: MockQuerier<Empty> = MockQuerier::new(&[(&addr, &balance)]);
125
126        let gas_limit = 20;
127        let (result, _gas_info) = querier.query_raw(b"broken request", gas_limit);
128        match result.unwrap_err() {
129            BackendError::OutOfGas {} => {}
130            err => panic!("Unexpected error: {err:?}"),
131        }
132    }
133
134    #[test]
135    fn bank_querier_balance() {
136        let addr = String::from("foobar");
137        let balance = vec![coin(123, "ELF"), coin(777, "FLY")];
138        let querier = MockQuerier::new(&[(&addr, &balance)]);
139
140        // one match
141        let fly = querier
142            .query::<Empty>(
143                &BankQuery::Balance {
144                    address: addr.clone(),
145                    denom: "FLY".to_string(),
146                }
147                .into(),
148                DEFAULT_QUERY_GAS_LIMIT,
149            )
150            .0
151            .unwrap()
152            .unwrap()
153            .unwrap();
154        let res: BalanceResponse = from_json(fly).unwrap();
155        assert_eq!(res.amount, coin(777, "FLY"));
156
157        // missing denom
158        let miss = querier
159            .query::<Empty>(
160                &BankQuery::Balance {
161                    address: addr,
162                    denom: "MISS".to_string(),
163                }
164                .into(),
165                DEFAULT_QUERY_GAS_LIMIT,
166            )
167            .0
168            .unwrap()
169            .unwrap()
170            .unwrap();
171        let res: BalanceResponse = from_json(miss).unwrap();
172        assert_eq!(res.amount, coin(0, "MISS"));
173    }
174
175    #[test]
176    #[allow(deprecated)]
177    fn bank_querier_missing_account() {
178        let addr = String::from("foobar");
179        let balance = vec![coin(123, "ELF"), coin(777, "FLY")];
180        let querier = MockQuerier::new(&[(&addr, &balance)]);
181
182        // any denom on balances on empty account is empty coin
183        let miss = querier
184            .query::<Empty>(
185                &BankQuery::Balance {
186                    address: String::from("elsewhere"),
187                    denom: "ELF".to_string(),
188                }
189                .into(),
190                DEFAULT_QUERY_GAS_LIMIT,
191            )
192            .0
193            .unwrap()
194            .unwrap()
195            .unwrap();
196        let res: BalanceResponse = from_json(miss).unwrap();
197        assert_eq!(res.amount, coin(0, "ELF"));
198    }
199}