Add option to use a denom list for distribution

public-awesome · Apr 26, 2024 · a7c312e · a7c312e
1 parent 1be018b
commit a7c312e
Show file tree

Hide file tree

Showing 4 changed files with 31 additions and 16 deletions.
diff --git a/contracts/splits/src/contract.rs b/contracts/splits/src/contract.rs
@@ -1,7 +1,7 @@
 #[cfg(not(feature = "library"))]
 use cosmwasm_std::entry_point;
 use cosmwasm_std::{
-    coins, ensure, to_json_binary, Addr, BankMsg, Binary, CosmosMsg, Deps, DepsMut, Env,
+    coins, ensure, to_json_binary, Addr, BankMsg, Binary, Coin, CosmosMsg, Deps, DepsMut, Env,
     MessageInfo, Reply, Response, StdResult, SubMsg, Uint128,
 };
 use cw2::set_contract_version;
@@ -70,14 +70,17 @@ pub fn execute(
         ExecuteMsg::UpdateAdmin { admin } => {
             Ok(ADMIN.execute_update_admin(deps, info, maybe_addr(api, admin)?)?)
         }
-        ExecuteMsg::Distribute {} => execute_distribute(deps.as_ref(), env, info),
+        ExecuteMsg::Distribute { denom_list } => {
+            execute_distribute(deps.as_ref(), env, info, denom_list)
+        }
     }
 }
 
 pub fn execute_distribute(
     deps: Deps,
     env: Env,
     info: MessageInfo,
+    denom_list: Option<Vec<String>>,
 ) -> Result<Response, ContractError> {
     if !can_distribute(deps, info)? {
         return Err(ContractError::Unauthorized {});
@@ -93,8 +96,20 @@ pub fn execute_distribute(
             count: members_count,
         });
     }
-
-    let funds = deps.querier.query_all_balances(env.contract.address)?;
+    let mut funds: Vec<Coin> = Vec::new();
+    if let Some(denom_list) = denom_list {
+        for denom in denom_list.iter() {
+            let balance = deps
+                .querier
+                .query_balance(env.contract.address.clone(), denom)?;
+            if balance.amount.is_zero() {
+                continue;
+            }
+            funds.push(balance);
+        }
+    } else {
+        funds = deps.querier.query_all_balances(env.contract.address)?;
+    }
 
     ensure!(!funds.is_empty(), ContractError::NoFunds {});
 

diff --git a/contracts/splits/src/msg.rs b/contracts/splits/src/msg.rs
@@ -17,7 +17,7 @@ pub struct InstantiateMsg {
 #[cw_serde]
 pub enum ExecuteMsg {
     UpdateAdmin { admin: Option<String> },
-    Distribute {},
+    Distribute { denom_list: Option<Vec<String>> },
 }
 
 #[cw_serde]

diff --git a/test-suite/src/splits/tests/integration_tests.rs b/test-suite/src/splits/tests/integration_tests.rs
@@ -270,7 +270,7 @@ mod tests {
 
             let (splits_addr, _) = setup_test_case(&mut app, vec![], false);
 
-            let msg = ExecuteMsg::Distribute {};
+            let msg = ExecuteMsg::Distribute { denom_list: None };
 
             let err = app
                 .execute_contract(Addr::unchecked(OWNER), splits_addr, &msg, &[])
@@ -286,10 +286,10 @@ mod tests {
 
             let (splits_addr, _) = setup_test_case(&mut app, init_funds, false);
 
-            let msg = ExecuteMsg::Distribute {};
+            let msg = ExecuteMsg::Distribute { denom_list: None };
 
             app.execute_contract(
-                Addr::unchecked("non_memeber".to_string()),
+                Addr::unchecked("non_member".to_string()),
                 splits_addr,
                 &msg,
                 &[],
@@ -305,7 +305,7 @@ mod tests {
 
             let (splits_addr, _) = setup_test_case_with_internal_group(&mut app, init_funds);
 
-            let msg = ExecuteMsg::Distribute {};
+            let msg = ExecuteMsg::Distribute { denom_list: None };
 
             app.execute_contract(Addr::unchecked(OWNER), splits_addr.clone(), &msg, &[])
                 .unwrap();
@@ -345,7 +345,7 @@ mod tests {
                 setup_test_case_with_internal_group(&mut app, init_funds);
             let total_weight = Cw4Contract(group_addr).total_weight(&app.wrap()).unwrap();
 
-            let msg = ExecuteMsg::Distribute {};
+            let msg = ExecuteMsg::Distribute { denom_list: None };
 
             let err = app
                 .execute_contract(Addr::unchecked(OWNER), splits_addr, &msg, &[])
@@ -369,7 +369,7 @@ mod tests {
             let multiplier = init_funds[0].amount / Uint128::from(total_weight);
             let contract_balance = init_funds[0].amount - multiplier * Uint128::from(total_weight);
 
-            let msg = ExecuteMsg::Distribute {};
+            let msg = ExecuteMsg::Distribute { denom_list: None };
 
             let _ = app
                 .execute_contract(Addr::unchecked(OWNER), splits_addr.clone(), &msg, &[])
@@ -403,7 +403,7 @@ mod tests {
 
             let (splits_addr, _) = setup_test_case_with_overflow_group(&mut app, init_funds);
 
-            let msg = ExecuteMsg::Distribute {};
+            let msg = ExecuteMsg::Distribute { denom_list: None };
             let err = app
                 .execute_contract(Addr::unchecked(OWNER), splits_addr, &msg, &[])
                 .unwrap_err();
@@ -433,7 +433,7 @@ mod tests {
                 .execute_contract(Addr::unchecked(OWNER), group_addr, &msg, &[])
                 .unwrap();
 
-            let msg = ExecuteMsg::Distribute {};
+            let msg = ExecuteMsg::Distribute { denom_list: None };
             let _ = app
                 .execute_contract(Addr::unchecked(OWNER), splits_addr, &msg, &[])
                 .unwrap();
@@ -465,7 +465,7 @@ mod tests {
             let contract_balance = init_funds[0].amount - multiplier * Uint128::from(total_weight);
             let mut payouts = vec![];
 
-            let msg = ExecuteMsg::Distribute {};
+            let msg = ExecuteMsg::Distribute { denom_list: None };
 
             let _ = app
                 .execute_contract(Addr::unchecked(OWNER), splits_addr.clone(), &msg, &[])
@@ -534,7 +534,7 @@ mod tests {
             );
 
             // distribute again and check accounting
-            let msg = ExecuteMsg::Distribute {};
+            let msg = ExecuteMsg::Distribute { denom_list: None };
             let _ = app
                 .execute_contract(Addr::unchecked(OWNER), splits_addr.clone(), &msg, &[])
                 .unwrap();

diff --git a/test-suite/src/vending_minter/tests/splits.rs b/test-suite/src/vending_minter/tests/splits.rs
@@ -95,7 +95,7 @@ fn mint_and_split() {
     );
     assert!(res.is_ok());
 
-    let dist_msg = SplitsExecuteMsg::Distribute {};
+    let dist_msg = SplitsExecuteMsg::Distribute { denom_list: None };
     let res = app.execute_contract(Addr::unchecked(OWNER), splits_addr, &dist_msg, &[]);
     assert!(res.is_ok());