From 83d60f7dff4e33053493715d66d10894b5c2228f Mon Sep 17 00:00:00 2001
From: Mykhailo Donchenko <91957742+Buckram123@users.noreply.github.com>
Date: Fri, 14 Jul 2023 15:31:03 +0300
Subject: [PATCH] Feat/Add balance function for Mock (#157)

* add_balance function

* updated changelog
---
 CHANGELOG.md             |  2 ++
 Cargo.toml               |  1 +
 cw-orch/Cargo.toml       |  1 +
 cw-orch/src/mock/core.rs | 42 ++++++++++++++++++++++++++++++++++++++++
 4 files changed, 46 insertions(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index d05bccdc9..18d26b704 100644
--- a/CHANGELOG.md
+++ b/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.
diff --git a/Cargo.toml b/Cargo.toml
index 9c5133eee..ac7fdeda6 100644
--- a/Cargo.toml
+++ b/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"
diff --git a/cw-orch/Cargo.toml b/cw-orch/Cargo.toml
index 73dd1b912..726e47e24 100644
--- a/cw-orch/Cargo.toml
+++ b/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 }
diff --git a/cw-orch/src/mock/core.rs b/cw-orch/src/mock/core.rs
index 2fbf02866..95dbe74d3 100644
--- a/cw-orch/src/mock/core.rs
+++ b/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)])
+    }
 }