Repository: DA0-DA0/dao-contracts (rust)

Issue:
Support DAO initial balance in staked balance voting module.
The old contracts support allocating a token amount to the DAO on instantiation. We should support this in the `TokenInfo::New` variant.

Pull request description:
Add DAO initial balance to staked governance token instantiation
Resolves #208 

Sending part of the initial token supply to the DAO was left out of the contract v1 rewrite. This adds it back in.

Patch:
diff --git a/contracts/cw20-staked-balance-voting/schema/instantiate_msg.json b/contracts/cw20-staked-balance-voting/schema/instantiate_msg.json
index c3048d25..b508d6ec 100644
--- a/contracts/cw20-staked-balance-voting/schema/instantiate_msg.json
+++ b/contracts/cw20-staked-balance-voting/schema/instantiate_msg.json
@@ -6,6 +6,16 @@
     "token_info"
   ],
   "properties": {
+    "initial_dao_balance": {
+      "anyOf": [
+        {
+          "$ref": "#/definitions/Uint128"
+        },
+        {
+          "type": "null"
+        }
+      ]
+    },
     "token_info": {
       "$ref": "#/definitions/TokenInfo"
     }
diff --git a/contracts/cw20-staked-balance-voting/schema/query_msg.json b/contracts/cw20-staked-balance-voting/schema/query_msg.json
index 3b351671..7e950e97 100644
--- a/contracts/cw20-staked-balance-voting/schema/query_msg.json
+++ b/contracts/cw20-staked-balance-voting/schema/query_msg.json
@@ -5,10 +5,10 @@
     {
       "type": "object",
       "required": [
-        "token_contract"
+        "staking_contract"
       ],
       "properties": {
-        "token_contract": {
+        "staking_contract": {
           "type": "object"
         }
       },
@@ -17,10 +17,10 @@
     {
       "type": "object",
       "required": [
-        "staking_contract"
+        "dao"
       ],
       "properties": {
-        "staking_contract": {
+        "dao": {
           "type": "object"
         }
       },
@@ -87,6 +87,18 @@
         }
       },
       "additionalProperties": false
+    },
+    {
+      "type": "object",
+      "required": [
+        "token_contract"
+      ],
+      "properties": {
+        "token_contract": {
+          "type": "object"
+        }
+      },
+      "additionalProperties": false
     }
   ]
 }
diff --git a/contracts/cw20-staked-balance-voting/src/contract.rs b/contracts/cw20-staked-balance-voting/src/contract.rs
index 31a71508..6b926572 100644
--- a/contracts/cw20-staked-balance-voting/src/contract.rs
+++ b/contracts/cw20-staked-balance-voting/src/contract.rs
@@ -5,6 +5,7 @@ use cosmwasm_std::{
     Uint128, WasmMsg,
 };
 use cw2::set_contract_version;
+use cw20::Cw20Coin;
 use cw_utils::parse_reply_instantiate_data;
 
 use crate::error::ContractError;
@@ -89,11 +90,21 @@ pub fn instantiate(
             name,
             symbol,
             decimals,
-            initial_balances,
+            mut initial_balances,
             marketing,
             staking_code_id,
             unstaking_duration,
         } => {
+            // Add DAO initial balance to initial_balances vector if defined.
+            if let Some(initial_dao_balance) = msg.initial_dao_balance {
+                if initial_dao_balance > Uint128::zero() {
+                    initial_balances.push(Cw20Coin {
+                        address: info.sender.to_string(),
+                        amount: initial_dao_balance,
+                    });
+                }
+            }
+
             let initial_supply = initial_balances
                 .iter()
                 .fold(Uint128::zero(), |p, n| p + n.amount);
diff --git a/contracts/cw20-staked-balance-voting/src/msg.rs b/contracts/cw20-staked-balance-voting/src/msg.rs
index e7c32def..d080f0d1 100644
--- a/contracts/cw20-staked-balance-voting/src/msg.rs
+++ b/contracts/cw20-staked-balance-voting/src/msg.rs
@@ -1,3 +1,4 @@
+use cosmwasm_std::Uint128;
 use cw20::Cw20Coin;
 use cw20_base::msg::InstantiateMarketingInfo;
 use cw_governance_macros::{token_query, voting_query};
@@ -41,6 +42,7 @@ pub enum TokenInfo {
 #[derive(Serialize, Deserialize, Clone, Debug, PartialEq, JsonSchema)]
 pub struct InstantiateMsg {
     pub token_info: TokenInfo,
+    pub initial_dao_balance: Option<Uint128>,
 }
 
 #[derive(Serialize, Deserialize, Clone, Debug, PartialEq, JsonSchema)]


Test patch:
diff --git a/contracts/cw20-staked-balance-voting/src/tests.rs b/contracts/cw20-staked-balance-voting/src/tests.rs
index c66c71f7..b7e05fde 100644
--- a/contracts/cw20-staked-balance-voting/src/tests.rs
+++ b/contracts/cw20-staked-balance-voting/src/tests.rs
@@ -1,6 +1,6 @@
 use cosmwasm_std::{to_binary, Addr, Empty, Uint128};
 use cw2::ContractVersion;
-use cw20::{Cw20Coin, MinterResponse, TokenInfoResponse};
+use cw20::{BalanceResponse, Cw20Coin, MinterResponse, TokenInfoResponse};
 use cw_governance_interface::voting::{InfoResponse, VotingPowerAtHeightResponse};
 use cw_multi_test::{next_block, App, Contract, ContractWrapper, Executor};
 
@@ -84,6 +84,7 @@ fn test_instantiate_zero_supply() {
                 unstaking_duration: None,
                 staking_code_id: staking_contract_id,
             },
+            initial_dao_balance: Some(Uint128::zero()),
         },
     );
 }
@@ -110,6 +111,7 @@ fn test_instantiate_no_balances() {
                 unstaking_duration: None,
                 staking_code_id: staking_contract_id,
             },
+            initial_dao_balance: Some(Uint128::zero()),
         },
     );
 }
@@ -139,6 +141,7 @@ fn test_contract_info() {
                 unstaking_duration: None,
                 staking_code_id: staking_contract_id,
             },
+            initial_dao_balance: Some(Uint128::zero()),
         },
     );
 
@@ -188,6 +191,7 @@ fn test_new_cw20() {
                 unstaking_duration: None,
                 staking_code_id: staking_contract_id,
             },
+            initial_dao_balance: Some(Uint128::from(10u64)),
         },
     );
 
@@ -210,7 +214,7 @@ fn test_new_cw20() {
             name: "DAO DAO".to_string(),
             symbol: "DAO".to_string(),
             decimals: 6,
-            total_supply: Uint128::from(2u64)
+            total_supply: Uint128::from(12u64)
         }
     );
 
@@ -226,6 +230,23 @@ fn test_new_cw20() {
         })
     );
 
+    // Expect DAO (sender address) to have initial balance.
+    let token_info: BalanceResponse = app
+        .wrap()
+        .query_wasm_smart(
+            token_addr.clone(),
+            &cw20::Cw20QueryMsg::Balance {
+                address: DAO_ADDR.to_string(),
+            },
+        )
+        .unwrap();
+    assert_eq!(
+        token_info,
+        BalanceResponse {
+            balance: Uint128::from(10u64)
+        }
+    );
+
     // Expect 0 as they have not staked
     let creator_voting_power: VotingPowerAtHeightResponse = app
         .wrap()
@@ -344,6 +365,7 @@ fn test_existing_cw20_new_staking() {
                     unstaking_duration: None,
                 },
             },
+            initial_dao_balance: None,
         },
     );
 
@@ -494,6 +516,7 @@ fn test_existing_cw20_existing_staking() {
                     unstaking_duration: None,
                 },
             },
+            initial_dao_balance: None,
         },
     );
 
@@ -531,6 +554,7 @@ fn test_existing_cw20_existing_staking() {
                     staking_contract_address: staking_addr.to_string(),
                 },
             },
+            initial_dao_balance: None,
         },
     );
 
@@ -645,6 +669,7 @@ fn test_existing_cw20_existing_staking() {
                     staking_contract_address: staking_addr.to_string(),
                 },
             },
+            initial_dao_balance: None,
         },
         &[],
         "voting module",
@@ -692,6 +717,7 @@ fn test_different_heights() {
                     unstaking_duration: None,
                 },
             },
+            initial_dao_balance: None,
         },
     );