cw-asset 4.0.0

Helper library for interacting with Cosmos assets (native coins and CW20 tokens)
Documentation
use std::collections::HashMap;

use cosmwasm_std::{to_json_binary, Addr, QuerierResult, SystemError, Uint128};
use cw20::{BalanceResponse, Cw20QueryMsg};

#[derive(Default)]
pub struct Cw20Querier {
    balances: HashMap<Addr, HashMap<Addr, Uint128>>,
}

impl Cw20Querier {
    pub fn handle_query(&self, contract_addr: &Addr, query: Cw20QueryMsg) -> QuerierResult {
        match query {
            Cw20QueryMsg::Balance {
                address,
            } => {
                let contract_balances = match self.balances.get(contract_addr) {
                    Some(balances) => balances,
                    None => {
                        return Err(SystemError::InvalidRequest {
                            error: format!(
                                "[mock]: cw20 balances not set for token {contract_addr:?}",
                            ),
                            request: Default::default(),
                        })
                        .into()
                    },
                };

                let balance = match contract_balances.get(&Addr::unchecked(&address)) {
                    Some(balance) => balance,
                    None => {
                        return Err(SystemError::InvalidRequest {
                            error: format!("[mock]: cw20 balance not set for user {address:?}"),
                            request: Default::default(),
                        })
                        .into()
                    },
                };

                Ok(to_json_binary(&BalanceResponse {
                    balance: *balance,
                })
                .into())
                .into()
            },

            query => Err(SystemError::InvalidRequest {
                error: format!("[mock]: unsupported cw20 query {query:?}"),
                request: Default::default(),
            })
            .into(),
        }
    }

    pub fn set_balance(&mut self, contract: &str, user: &str, balance: u128) {
        let contract_addr = Addr::unchecked(contract);
        let user_addr = Addr::unchecked(user);

        let contract_balances = self.balances.entry(contract_addr).or_default();
        contract_balances.insert(user_addr, Uint128::new(balance));
    }
}