terra_asset/testing/
cw20_querier.rs

1use cosmwasm_std::{to_binary, Addr, QuerierResult, SystemError, Uint128};
2use cw20::{BalanceResponse, Cw20QueryMsg};
3use std::collections::HashMap;
4
5#[derive(Default)]
6pub struct Cw20Querier {
7    balances: HashMap<Addr, HashMap<Addr, Uint128>>,
8}
9
10impl Cw20Querier {
11    pub fn handle_query(&self, contract_addr: &Addr, query: Cw20QueryMsg) -> QuerierResult {
12        match query {
13            Cw20QueryMsg::Balance {
14                address,
15            } => {
16                let contract_balances = match self.balances.get(contract_addr) {
17                    Some(balances) => balances,
18                    None => {
19                        return Err(SystemError::InvalidRequest {
20                            error: format!(
21                                "[mock]: cw20 balances not set for token {:?}",
22                                contract_addr
23                            ),
24                            request: Default::default(),
25                        })
26                        .into()
27                    }
28                };
29
30                let balance = match contract_balances.get(&Addr::unchecked(&address)) {
31                    Some(balance) => balance,
32                    None => {
33                        return Err(SystemError::InvalidRequest {
34                            error: format!("[mock]: cw20 balance not set for user {:?}", address),
35                            request: Default::default(),
36                        })
37                        .into()
38                    }
39                };
40
41                Ok(to_binary(&BalanceResponse {
42                    balance: *balance,
43                })
44                .into())
45                .into()
46            }
47
48            query => Err(SystemError::InvalidRequest {
49                error: format!("[mock]: unsupported cw20 query {:?}", query),
50                request: Default::default(),
51            })
52            .into(),
53        }
54    }
55
56    pub fn set_balance(&mut self, contract: &str, user: &str, balance: u128) {
57        let contract_addr = Addr::unchecked(contract);
58        let user_addr = Addr::unchecked(user);
59
60        let contract_balances = self.balances.entry(contract_addr).or_insert_with(HashMap::new);
61        contract_balances.insert(user_addr, Uint128::new(balance));
62    }
63}
terra_asset/testing/cw20_querier.rs

terra_asset/testing/
cw20_querier.rs
