1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137

use std::cell::RefCell;
use std::rc::Rc;

use cosmwasm_std::testing::MockApi;
use cosmwasm_std::{Addr, Coin, Uint128};
use cw_multi_test::AppBuilder;
use cw_orch_core::environment::{BankQuerier, BankSetter, TxHandler};
use cw_orch_core::{
    environment::{DefaultQueriers, StateInterface},
    CwEnvError,
};
use cw_utils::NativeBalance;

use crate::queriers::bank::MockBankQuerier;
use crate::{Mock, MockState};

impl<S: StateInterface> Mock<S> {
    /// Set the bank balance of an address.
    pub fn set_balance(
        &self,
        address: impl Into<String>,
        amount: Vec<cosmwasm_std::Coin>,
    ) -> Result<(), CwEnvError> {
        self.app
            .borrow_mut()
            .init_modules(|router, _, storage| {
                router
                    .bank
                    .init_balance(storage, &Addr::unchecked(address.into()), amount)
            })
            .map_err(Into::into)
    }

    /// Adds the bank balance of an address.
    pub fn add_balance(
        &self,
        address: impl Into<String>,
        amount: Vec<cosmwasm_std::Coin>,
    ) -> Result<(), CwEnvError> {
        let addr = &Addr::unchecked(address.into());
        let b = self.query_all_balances(addr.clone())?;
        let new_amount = NativeBalance(b) + NativeBalance(amount);
        self.app
            .borrow_mut()
            .init_modules(|router, _, storage| {
                router
                    .bank
                    .init_balance(storage, addr, new_amount.into_vec())
            })
            .map_err(Into::into)
    }

    /// Set the balance for multiple coins at once.
    pub fn set_balances(
        &self,
        balances: &[(impl Into<String> + Clone, &[cosmwasm_std::Coin])],
    ) -> Result<(), CwEnvError> {
        self.app
            .borrow_mut()
            .init_modules(|router, _, storage| -> Result<(), CwEnvError> {
                for (addr, coins) in balances {
                    router.bank.init_balance(
                        storage,
                        &Addr::unchecked(addr.clone()),
                        coins.to_vec(),
                    )?;
                }
                Ok(())
            })
    }

    /// Query the (bank) balance of a native token for and address.
    /// Returns the amount of the native token.
    pub fn query_balance(
        &self,
        address: impl Into<String>,
        denom: &str,
    ) -> Result<Uint128, CwEnvError> {
        Ok(self
            .bank_querier()
            .balance(address, Some(denom.to_string()))?
            .first()
            .map(|c| c.amount)
            .unwrap_or_default())
    }

    /// Fetch all the balances of an address.
    pub fn query_all_balances(
        &self,
        address: impl Into<String>,
    ) -> Result<Vec<cosmwasm_std::Coin>, CwEnvError> {
        self.bank_querier().balance(address, None)
    }
}

impl Mock {
    /// Create a mock environment with the default mock state.
    pub fn new(sender: impl Into<String>) -> Self {
        Mock::new_custom(sender, MockState::new())
    }

    pub fn new_with_chain_id(sender: impl Into<String>, chain_id: &str) -> Self {
        let chain = Mock::new_custom(sender, MockState::new());
        chain
            .app
            .borrow_mut()
            .update_block(|b| b.chain_id = chain_id.to_string());

        chain
    }
}
impl<S: StateInterface> Mock<S> {
    /// Create a mock environment with a custom mock state.
    /// The state is customizable by implementing the `StateInterface` trait on a custom struct and providing it on the custom constructor.
    pub fn new_custom(sender: impl Into<String>, custom_state: S) -> Self {
        let state = Rc::new(RefCell::new(custom_state));
        let app = Rc::new(RefCell::new(AppBuilder::new_custom().build(|_, _, _| {})));

        Self {
            sender: Addr::unchecked(sender),
            state,
            app,
        }
    }
}

impl<S: StateInterface> BankSetter for Mock<S> {
    type T = MockBankQuerier<MockApi>;

    fn set_balance(
        &mut self,
        address: impl Into<String>,
        amount: Vec<Coin>,
    ) -> Result<(), <Self as TxHandler>::Error> {
        (*self).set_balance(address, amount)
    }
}