apollo_cw_multi_test/
bank.rs

1use anyhow::{bail, Result as AnyResult};
2use itertools::Itertools;
3
4use schemars::JsonSchema;
5
6use cosmwasm_std::{
7    coin, to_binary, Addr, AllBalanceResponse, Api, BalanceResponse, BankMsg, BankQuery, Binary,
8    BlockInfo, Coin, Event, Order, Querier, StdResult, Storage, SupplyResponse, Uint128,
9};
10use cw_storage_plus::Map;
11use cw_utils::NativeBalance;
12
13use crate::app::CosmosRouter;
14use crate::executor::AppResponse;
15use crate::module::Module;
16use crate::prefixed_storage::{prefixed, prefixed_read};
17
18const BALANCES: Map<&Addr, NativeBalance> = Map::new("balances");
19
20pub const NAMESPACE_BANK: &[u8] = b"bank";
21
22// WIP
23#[derive(Clone, std::fmt::Debug, PartialEq, Eq, JsonSchema)]
24pub enum BankSudo {
25    Mint {
26        to_address: String,
27        amount: Vec<Coin>,
28    },
29}
30
31pub trait Bank: Module<ExecT = BankMsg, QueryT = BankQuery, SudoT = BankSudo> {}
32
33#[derive(Default, Clone)]
34pub struct BankKeeper {}
35
36impl BankKeeper {
37    pub fn new() -> Self {
38        BankKeeper {}
39    }
40
41    // this is an "admin" function to let us adjust bank accounts in genesis
42    pub fn init_balance(
43        &self,
44        storage: &mut dyn Storage,
45        account: &Addr,
46        amount: Vec<Coin>,
47    ) -> AnyResult<()> {
48        let mut bank_storage = prefixed(storage, NAMESPACE_BANK);
49        self.set_balance(&mut bank_storage, account, amount)
50    }
51
52    fn set_balance(
53        &self,
54        bank_storage: &mut dyn Storage,
55        account: &Addr,
56        amount: Vec<Coin>,
57    ) -> AnyResult<()> {
58        let mut balance = NativeBalance(amount);
59        balance.normalize();
60        BALANCES
61            .save(bank_storage, account, &balance)
62            .map_err(Into::into)
63    }
64
65    // this is an "admin" function to let us adjust bank accounts
66    fn get_balance(&self, bank_storage: &dyn Storage, account: &Addr) -> AnyResult<Vec<Coin>> {
67        let val = BALANCES.may_load(bank_storage, account)?;
68        Ok(val.unwrap_or_default().into_vec())
69    }
70
71    fn get_supply(&self, bank_storage: &dyn Storage, denom: &str) -> AnyResult<Uint128> {
72        Ok(BALANCES
73            .range(bank_storage, None, None, Order::Ascending)
74            .map(|item| item.map(|i| i.1.into_vec()))
75            .collect::<StdResult<Vec<Vec<Coin>>>>()?
76            .into_iter()
77            .flatten()
78            .filter(|c| &c.denom == denom)
79            .map(|c| c.amount)
80            .sum::<Uint128>())
81    }
82
83    fn send(
84        &self,
85        bank_storage: &mut dyn Storage,
86        from_address: Addr,
87        to_address: Addr,
88        amount: Vec<Coin>,
89    ) -> AnyResult<()> {
90        self.burn(bank_storage, from_address, amount.clone())?;
91        self.mint(bank_storage, to_address, amount)
92    }
93
94    fn mint(
95        &self,
96        bank_storage: &mut dyn Storage,
97        to_address: Addr,
98        amount: Vec<Coin>,
99    ) -> AnyResult<()> {
100        let amount = self.normalize_amount(amount)?;
101        let b = self.get_balance(bank_storage, &to_address)?;
102        let b = NativeBalance(b) + NativeBalance(amount);
103        self.set_balance(bank_storage, &to_address, b.into_vec())
104    }
105
106    fn burn(
107        &self,
108        bank_storage: &mut dyn Storage,
109        from_address: Addr,
110        amount: Vec<Coin>,
111    ) -> AnyResult<()> {
112        let amount = self.normalize_amount(amount)?;
113        let a = self.get_balance(bank_storage, &from_address)?;
114        let a = (NativeBalance(a) - amount)?;
115        self.set_balance(bank_storage, &from_address, a.into_vec())
116    }
117
118    /// Filters out all 0 value coins and returns an error if the resulting Vec is empty
119    fn normalize_amount(&self, amount: Vec<Coin>) -> AnyResult<Vec<Coin>> {
120        let res: Vec<_> = amount.into_iter().filter(|x| !x.amount.is_zero()).collect();
121        if res.is_empty() {
122            bail!("Cannot transfer empty coins amount")
123        } else {
124            Ok(res)
125        }
126    }
127}
128
129fn coins_to_string(coins: &[Coin]) -> String {
130    coins
131        .iter()
132        .map(|c| format!("{}{}", c.amount, c.denom))
133        .join(",")
134}
135
136impl Bank for BankKeeper {}
137
138impl Module for BankKeeper {
139    type ExecT = BankMsg;
140    type QueryT = BankQuery;
141    type SudoT = BankSudo;
142
143    fn execute<ExecC, QueryC>(
144        &self,
145        _api: &dyn Api,
146        storage: &mut dyn Storage,
147        _router: &dyn CosmosRouter<ExecC = ExecC, QueryC = QueryC>,
148        _block: &BlockInfo,
149        sender: Addr,
150        msg: BankMsg,
151    ) -> AnyResult<AppResponse> {
152        let mut bank_storage = prefixed(storage, NAMESPACE_BANK);
153        match msg {
154            BankMsg::Send { to_address, amount } => {
155                // see https://github.com/cosmos/cosmos-sdk/blob/v0.42.7/x/bank/keeper/send.go#L142-L147
156                let events = vec![Event::new("transfer")
157                    .add_attribute("recipient", &to_address)
158                    .add_attribute("sender", &sender)
159                    .add_attribute("amount", coins_to_string(&amount))];
160                self.send(
161                    &mut bank_storage,
162                    sender,
163                    Addr::unchecked(to_address),
164                    amount,
165                )?;
166                Ok(AppResponse { events, data: None })
167            }
168            BankMsg::Burn { amount } => {
169                // burn doesn't seem to emit any events
170                self.burn(&mut bank_storage, sender, amount)?;
171                Ok(AppResponse::default())
172            }
173            m => bail!("Unsupported bank message: {:?}", m),
174        }
175    }
176
177    fn sudo<ExecC, QueryC>(
178        &self,
179        api: &dyn Api,
180        storage: &mut dyn Storage,
181        _router: &dyn CosmosRouter<ExecC = ExecC, QueryC = QueryC>,
182        _block: &BlockInfo,
183        msg: BankSudo,
184    ) -> AnyResult<AppResponse> {
185        let mut bank_storage = prefixed(storage, NAMESPACE_BANK);
186        match msg {
187            BankSudo::Mint { to_address, amount } => {
188                let to_address = api.addr_validate(&to_address)?;
189                self.mint(&mut bank_storage, to_address, amount)?;
190                Ok(AppResponse::default())
191            }
192        }
193    }
194
195    fn query(
196        &self,
197        api: &dyn Api,
198        storage: &dyn Storage,
199        _querier: &dyn Querier,
200        _block: &BlockInfo,
201        request: BankQuery,
202    ) -> AnyResult<Binary> {
203        let bank_storage = prefixed_read(storage, NAMESPACE_BANK);
204        match request {
205            BankQuery::AllBalances { address } => {
206                let address = api.addr_validate(&address)?;
207                let amount = self.get_balance(&bank_storage, &address)?;
208                let res = AllBalanceResponse { amount };
209                Ok(to_binary(&res)?)
210            }
211            BankQuery::Balance { address, denom } => {
212                let address = api.addr_validate(&address)?;
213                let all_amounts = self.get_balance(&bank_storage, &address)?;
214                let amount = all_amounts
215                    .into_iter()
216                    .find(|c| c.denom == denom)
217                    .unwrap_or_else(|| coin(0, denom));
218                let res = BalanceResponse { amount };
219                Ok(to_binary(&res)?)
220            }
221            BankQuery::Supply { denom } => {
222                let supply = self.get_supply(&bank_storage, &denom)?;
223                let mut res = SupplyResponse::default();
224                res.amount = Coin::new(supply.u128(), denom);
225                Ok(to_binary(&res)?)
226            }
227            q => bail!("Unsupported bank query: {:?}", q),
228        }
229    }
230}
231
232// Temp solution to get the query service paths. TODO: Figure out how to get this from the proto files (PR to osmosis-rust?)
233
234#[cfg(test)]
235mod test {
236    use super::*;
237
238    use crate::app::MockRouter;
239    use cosmwasm_std::testing::{mock_env, MockApi, MockQuerier, MockStorage};
240    use cosmwasm_std::{coins, from_slice, Empty, StdError};
241
242    fn query_balance(
243        bank: &BankKeeper,
244        api: &dyn Api,
245        store: &dyn Storage,
246        rcpt: &Addr,
247    ) -> Vec<Coin> {
248        let req = BankQuery::AllBalances {
249            address: rcpt.clone().into(),
250        };
251        let block = mock_env().block;
252        let querier: MockQuerier<Empty> = MockQuerier::new(&[]);
253
254        let raw = bank.query(api, store, &querier, &block, req).unwrap();
255        let res: AllBalanceResponse = from_slice(&raw).unwrap();
256        res.amount
257    }
258
259    #[test]
260    fn get_set_balance() {
261        let api = MockApi::default();
262        let mut store = MockStorage::new();
263        let block = mock_env().block;
264        let querier: MockQuerier<Empty> = MockQuerier::new(&[]);
265
266        let owner = Addr::unchecked("owner");
267        let rcpt = Addr::unchecked("receiver");
268        let init_funds = vec![coin(100, "eth"), coin(20, "btc")];
269        let norm = vec![coin(20, "btc"), coin(100, "eth")];
270
271        // set money
272        let bank = BankKeeper::new();
273        bank.init_balance(&mut store, &owner, init_funds).unwrap();
274        let bank_storage = prefixed_read(&store, NAMESPACE_BANK);
275
276        // get balance work
277        let rich = bank.get_balance(&bank_storage, &owner).unwrap();
278        assert_eq!(rich, norm);
279        let poor = bank.get_balance(&bank_storage, &rcpt).unwrap();
280        assert_eq!(poor, vec![]);
281
282        // proper queries work
283        let req = BankQuery::AllBalances {
284            address: owner.clone().into(),
285        };
286        let raw = bank.query(&api, &store, &querier, &block, req).unwrap();
287        let res: AllBalanceResponse = from_slice(&raw).unwrap();
288        assert_eq!(res.amount, norm);
289
290        let req = BankQuery::AllBalances {
291            address: rcpt.clone().into(),
292        };
293        let raw = bank.query(&api, &store, &querier, &block, req).unwrap();
294        let res: AllBalanceResponse = from_slice(&raw).unwrap();
295        assert_eq!(res.amount, vec![]);
296
297        let req = BankQuery::Balance {
298            address: owner.clone().into(),
299            denom: "eth".into(),
300        };
301        let raw = bank.query(&api, &store, &querier, &block, req).unwrap();
302        let res: BalanceResponse = from_slice(&raw).unwrap();
303        assert_eq!(res.amount, coin(100, "eth"));
304
305        let req = BankQuery::Balance {
306            address: owner.into(),
307            denom: "foobar".into(),
308        };
309        let raw = bank.query(&api, &store, &querier, &block, req).unwrap();
310        let res: BalanceResponse = from_slice(&raw).unwrap();
311        assert_eq!(res.amount, coin(0, "foobar"));
312
313        let req = BankQuery::Balance {
314            address: rcpt.into(),
315            denom: "eth".into(),
316        };
317        let raw = bank.query(&api, &store, &querier, &block, req).unwrap();
318        let res: BalanceResponse = from_slice(&raw).unwrap();
319        assert_eq!(res.amount, coin(0, "eth"));
320    }
321
322    #[test]
323    fn send_coins() {
324        let api = MockApi::default();
325        let mut store = MockStorage::new();
326        let block = mock_env().block;
327        let router = MockRouter::default();
328
329        let owner = Addr::unchecked("owner");
330        let rcpt = Addr::unchecked("receiver");
331        let init_funds = vec![coin(20, "btc"), coin(100, "eth")];
332        let rcpt_funds = vec![coin(5, "btc")];
333
334        // set money
335        let bank = BankKeeper::new();
336        bank.init_balance(&mut store, &owner, init_funds).unwrap();
337        bank.init_balance(&mut store, &rcpt, rcpt_funds).unwrap();
338
339        // send both tokens
340        let to_send = vec![coin(30, "eth"), coin(5, "btc")];
341        let msg = BankMsg::Send {
342            to_address: rcpt.clone().into(),
343            amount: to_send,
344        };
345        bank.execute(
346            &api,
347            &mut store,
348            &router,
349            &block,
350            owner.clone(),
351            msg.clone(),
352        )
353        .unwrap();
354        let rich = query_balance(&bank, &api, &store, &owner);
355        assert_eq!(vec![coin(15, "btc"), coin(70, "eth")], rich);
356        let poor = query_balance(&bank, &api, &store, &rcpt);
357        assert_eq!(vec![coin(10, "btc"), coin(30, "eth")], poor);
358
359        // can send from any account with funds
360        bank.execute(&api, &mut store, &router, &block, rcpt.clone(), msg)
361            .unwrap();
362
363        // cannot send too much
364        let msg = BankMsg::Send {
365            to_address: rcpt.into(),
366            amount: coins(20, "btc"),
367        };
368        bank.execute(&api, &mut store, &router, &block, owner.clone(), msg)
369            .unwrap_err();
370
371        let rich = query_balance(&bank, &api, &store, &owner);
372        assert_eq!(vec![coin(15, "btc"), coin(70, "eth")], rich);
373    }
374
375    #[test]
376    fn burn_coins() {
377        let api = MockApi::default();
378        let mut store = MockStorage::new();
379        let block = mock_env().block;
380        let router = MockRouter::default();
381
382        let owner = Addr::unchecked("owner");
383        let rcpt = Addr::unchecked("recipient");
384        let init_funds = vec![coin(20, "btc"), coin(100, "eth")];
385
386        // set money
387        let bank = BankKeeper::new();
388        bank.init_balance(&mut store, &owner, init_funds).unwrap();
389
390        // burn both tokens
391        let to_burn = vec![coin(30, "eth"), coin(5, "btc")];
392        let msg = BankMsg::Burn { amount: to_burn };
393        bank.execute(&api, &mut store, &router, &block, owner.clone(), msg)
394            .unwrap();
395        let rich = query_balance(&bank, &api, &store, &owner);
396        assert_eq!(vec![coin(15, "btc"), coin(70, "eth")], rich);
397
398        // cannot burn too much
399        let msg = BankMsg::Burn {
400            amount: coins(20, "btc"),
401        };
402        let err = bank
403            .execute(&api, &mut store, &router, &block, owner.clone(), msg)
404            .unwrap_err();
405        assert!(matches!(err.downcast().unwrap(), StdError::Overflow { .. }));
406
407        let rich = query_balance(&bank, &api, &store, &owner);
408        assert_eq!(vec![coin(15, "btc"), coin(70, "eth")], rich);
409
410        // cannot burn from empty account
411        let msg = BankMsg::Burn {
412            amount: coins(1, "btc"),
413        };
414        let err = bank
415            .execute(&api, &mut store, &router, &block, rcpt, msg)
416            .unwrap_err();
417        assert!(matches!(err.downcast().unwrap(), StdError::Overflow { .. }));
418    }
419
420    #[test]
421    fn fail_on_zero_values() {
422        let api = MockApi::default();
423        let mut store = MockStorage::new();
424        let block = mock_env().block;
425        let router = MockRouter::default();
426
427        let owner = Addr::unchecked("owner");
428        let rcpt = Addr::unchecked("recipient");
429        let init_funds = vec![coin(5000, "atom"), coin(100, "eth")];
430
431        // set money
432        let bank = BankKeeper::new();
433        bank.init_balance(&mut store, &owner, init_funds).unwrap();
434
435        // can send normal amounts
436        let msg = BankMsg::Send {
437            to_address: rcpt.to_string(),
438            amount: coins(100, "atom"),
439        };
440        bank.execute(&api, &mut store, &router, &block, owner.clone(), msg)
441            .unwrap();
442
443        // fails send on no coins
444        let msg = BankMsg::Send {
445            to_address: rcpt.to_string(),
446            amount: vec![],
447        };
448        bank.execute(&api, &mut store, &router, &block, owner.clone(), msg)
449            .unwrap_err();
450
451        // fails send on 0 coins
452        let msg = BankMsg::Send {
453            to_address: rcpt.to_string(),
454            amount: coins(0, "atom"),
455        };
456        bank.execute(&api, &mut store, &router, &block, owner.clone(), msg)
457            .unwrap_err();
458
459        // fails burn on no coins
460        let msg = BankMsg::Burn { amount: vec![] };
461        bank.execute(&api, &mut store, &router, &block, owner.clone(), msg)
462            .unwrap_err();
463
464        // fails burn on 0 coins
465        let msg = BankMsg::Burn {
466            amount: coins(0, "atom"),
467        };
468        bank.execute(&api, &mut store, &router, &block, owner, msg)
469            .unwrap_err();
470
471        // can mint via sudo
472        let msg = BankSudo::Mint {
473            to_address: rcpt.to_string(),
474            amount: coins(4321, "atom"),
475        };
476        bank.sudo(&api, &mut store, &router, &block, msg).unwrap();
477
478        // mint fails with 0 tokens
479        let msg = BankSudo::Mint {
480            to_address: rcpt.to_string(),
481            amount: coins(0, "atom"),
482        };
483        bank.sudo(&api, &mut store, &router, &block, msg)
484            .unwrap_err();
485
486        // mint fails with no tokens
487        let msg = BankSudo::Mint {
488            to_address: rcpt.to_string(),
489            amount: vec![],
490        };
491        bank.sudo(&api, &mut store, &router, &block, msg)
492            .unwrap_err();
493    }
494
495    #[test]
496    fn query_total_supply() {
497        let api = MockApi::default();
498        let mut storage = MockStorage::new();
499        let block = mock_env().block;
500        let querier: MockQuerier<Empty> = MockQuerier::new(&[]);
501
502        let addr1 = Addr::unchecked("addr1");
503        let addr2 = Addr::unchecked("addr2");
504        let init_funds = vec![coin(5000, "atom"), coin(100, "eth")];
505
506        // Init balances
507        let bank = BankKeeper::new();
508        bank.init_balance(&mut storage, &addr1, init_funds.clone())
509            .unwrap();
510        bank.init_balance(&mut storage, &addr2, init_funds).unwrap();
511
512        // Query total supply of atom
513        let request = BankQuery::Supply {
514            denom: "atom".to_string(),
515        };
516        let raw = bank
517            .query(&api, &storage, &querier, &block, request)
518            .unwrap();
519        let supply_res: SupplyResponse = from_slice(&raw).unwrap();
520        assert_eq!(supply_res.amount, Coin::new(10000u128, "atom"));
521
522        // Query total supply of eth
523        let request = BankQuery::Supply {
524            denom: "eth".to_string(),
525        };
526        let raw = bank
527            .query(&api, &storage, &querier, &block, request)
528            .unwrap();
529        let supply_res: SupplyResponse = from_slice(&raw).unwrap();
530        assert_eq!(supply_res.amount, Coin::new(200u128, "eth"));
531
532        // Query total supply of btc
533        let request = BankQuery::Supply {
534            denom: "btc".to_string(),
535        };
536        let raw = bank
537            .query(&api, &storage, &querier, &block, request)
538            .unwrap();
539        let supply_res: SupplyResponse = from_slice(&raw).unwrap();
540        assert_eq!(supply_res.amount, Coin::new(0u128, "btc"));
541    }
542}
apollo_cw_multi_test/bank.rs

apollo_cw_multi_test/
bank.rs
