terra_asset/testing/
custom_mock_querier.rs

1use cosmwasm_std::testing::MockQuerier;
2use cosmwasm_std::{
3    from_binary, from_slice, Addr, Coin, Decimal, Querier, QuerierResult, QueryRequest, StdResult,
4    SystemError, WasmQuery,
5};
6use cw20::Cw20QueryMsg;
7use terra_cosmwasm::TerraQueryWrapper;
8
9use super::{cw20_querier::Cw20Querier, native_querier::NativeQuerier};
10
11pub struct CustomMockQuerier {
12    base: MockQuerier<TerraQueryWrapper>,
13    native_querier: NativeQuerier,
14    cw20_querier: Cw20Querier,
15}
16
17impl Default for CustomMockQuerier {
18    fn default() -> Self {
19        CustomMockQuerier {
20            base: MockQuerier::new(&[]),
21            native_querier: NativeQuerier::default(),
22            cw20_querier: Cw20Querier::default(),
23        }
24    }
25}
26
27impl Querier for CustomMockQuerier {
28    fn raw_query(&self, bin_request: &[u8]) -> QuerierResult {
29        let request: QueryRequest<TerraQueryWrapper> = match from_slice(bin_request) {
30            Ok(v) => v,
31            Err(e) => {
32                return Err(SystemError::InvalidRequest {
33                    error: format!("[mock]: failed to parse query request {}", e),
34                    request: bin_request.into(),
35                })
36                .into()
37            }
38        };
39        self.handle_query(&request)
40    }
41}
42
43impl CustomMockQuerier {
44    pub fn handle_query(&self, request: &QueryRequest<TerraQueryWrapper>) -> QuerierResult {
45        match request {
46            QueryRequest::Custom(TerraQueryWrapper {
47                route,
48                query_data,
49            }) => self.native_querier.handle_query(route, query_data),
50
51            QueryRequest::Wasm(WasmQuery::Smart {
52                contract_addr,
53                msg,
54            }) => {
55                let contract_addr = Addr::unchecked(contract_addr);
56
57                let parse_cw20_query: StdResult<Cw20QueryMsg> = from_binary(msg);
58                if let Ok(cw20_query) = parse_cw20_query {
59                    return self.cw20_querier.handle_query(&contract_addr, cw20_query);
60                }
61
62                panic!("[mock]: unsupported wasm query {:?}", msg);
63            }
64
65            _ => self.base.handle_query(request),
66        }
67    }
68
69    pub fn set_base_balances(&mut self, address: &str, balances: &[Coin]) {
70        self.base.update_balance(address, balances.to_vec());
71    }
72
73    pub fn set_cw20_balance(&mut self, contract: &str, user: &str, balance: u128) {
74        self.cw20_querier.set_balance(contract, user, balance);
75    }
76
77    pub fn set_native_tax_rate(&mut self, tax_rate: Decimal) {
78        self.native_querier.set_tax_rate(tax_rate);
79    }
80
81    pub fn set_native_tax_cap(&mut self, denom: &str, tax_cap: u128) {
82        self.native_querier.set_tax_cap(denom, tax_cap);
83    }
84}
terra_asset/testing/custom_mock_querier.rs

terra_asset/testing/
custom_mock_querier.rs
