Feat/Add balance function for Mock (#157)

* add_balance function

* updated changelog

CHANGELOG.md

@@ -2,6 +2,8 @@
 
 ## Unpublished
 
+- Add `add_balance` function on the `Mock` type. 
+
 ## v0.10.0
 
 - Update `CallAs` to be generic over environments.

Cargo.toml

@@ -10,6 +10,7 @@ license = "GPL-3.0-only"
 repository = "https://github.com/AbstractSDK/cw-orchestrator"
 
 [workspace.dependencies]
+cw-utils = { version = "1.0.1" }
 cosmwasm-std = { version = "1.1" }
 cw-multi-test = { version = "0.16.4" }
 anyhow = "1.0"

cw-orch/Cargo.toml

@@ -68,6 +68,7 @@ cw-orch-contract-derive = { path = "../packages/cw-orch-contract-derive", versio
 cw-orch-fns-derive = { path = "../packages/cw-orch-fns-derive", version = "0.13.3" }
 
 cosmwasm-std = { workspace = true }
+cw-utils = { workspace = true }
 cw-multi-test = { workspace = true }
 anyhow = { workspace = true }
 serde = { workspace = true }

cw-orch/src/mock/core.rs

@@ -2,6 +2,7 @@ use std::{cell::RefCell, fmt::Debug, rc::Rc};
 
 use cosmwasm_std::{Addr, Empty, Event, Uint128};
 use cw_multi_test::{custom_app, next_block, AppResponse, BasicApp, Contract, Executor};
+use cw_utils::NativeBalance;
 use serde::{de::DeserializeOwned, Serialize};
 
 use crate::{
@@ -68,6 +69,24 @@ impl<S: StateInterface> Mock<S> {
             .map_err(Into::into)
     }
 
+    /// Adds the bank balance of an address.
+    pub fn add_balance(
+        &self,
+        address: &Addr,
+        amount: Vec<cosmwasm_std::Coin>,
+    ) -> Result<(), CwOrchError> {
+        let b = self.query_all_balances(address)?;
+        let new_amount = NativeBalance(b) + NativeBalance(amount);
+        self.app
+            .borrow_mut()
+            .init_modules(|router, _, storage| {
+                router
+                    .bank
+                    .init_balance(storage, address, new_amount.into_vec())
+            })
+            .map_err(Into::into)
+    }
+
     /// Set the balance for multiple coins at once.
     pub fn set_balances(
         &self,
@@ -490,4 +509,27 @@ mod test {
             .that(&mock_state.get_all_addresses().unwrap().len())
             .is_equal_to(1);
     }
+
+    #[test]
+    fn add_balance() {
+        let sender = &Addr::unchecked(SENDER);
+        let recipient = &Addr::unchecked(BALANCE_ADDR);
+        let amount = 1000000u128;
+        let denom_1 = "uosmo";
+        let denom_2 = "osmou";
+
+        let chain = Mock::new(sender);
+
+        chain
+            .add_balance(recipient, vec![Coin::new(amount, denom_1)])
+            .unwrap();
+        chain
+            .add_balance(recipient, vec![Coin::new(amount, denom_2)])
+            .unwrap();
+
+        let balances = chain.query_all_balances(recipient).unwrap();
+        asserting("recipient balances added")
+            .that(&balances)
+            .contains_all_of(&[&Coin::new(amount, denom_1), &Coin::new(amount, denom_2)])
+    }
 }