From e9112540482f04ee5a8492c821363a2c3a37699f Mon Sep 17 00:00:00 2001
From: Tony Chen <codchen03@gmail.com>
Date: Mon, 2 Dec 2024 15:03:56 +0800
Subject: [PATCH] allow all states to be override in instantiation

---
 src/contract.rs | 21 ++++++++++++++++++---
 src/msg.rs      |  5 +++++
 2 files changed, 23 insertions(+), 3 deletions(-)

diff --git a/src/contract.rs b/src/contract.rs
index ef593a5..017060f 100644
--- a/src/contract.rs
+++ b/src/contract.rs
@@ -233,9 +233,12 @@ pub fn instantiate(
             percentage: Decimal::percent(msg.admin_voting_threshold_percentage as u64),
         },
     )?;
-    WITHDRAWN_STAKING_REWARDS.save(deps.storage, &0)?;
-    WITHDRAWN_UNLOCKED.save(deps.storage, &0)?;
-    WITHDRAWN_LOCKED.save(deps.storage, &0)?;
+    let withdrawn_rewards = msg.withdrawn_staking_rewards.unwrap_or(0);
+    WITHDRAWN_STAKING_REWARDS.save(deps.storage, &withdrawn_rewards)?;
+    let withdrawn_unlocked = msg.withdrawn_unlocked.unwrap_or(0);
+    WITHDRAWN_UNLOCKED.save(deps.storage, &withdrawn_unlocked)?;
+    let withdrawn_locked = msg.withdrawn_locked.unwrap_or(0);
+    WITHDRAWN_LOCKED.save(deps.storage, &withdrawn_locked)?;
     Ok(Response::default())
 }
 
@@ -850,6 +853,9 @@ mod tests {
             },
             max_voting_period: Duration::Time(3600),
             admin_voting_threshold_percentage: 75,
+            withdrawn_staking_rewards: None,
+            withdrawn_locked: None,
+            withdrawn_unlocked: None,
         };
         instantiate(deps, mock_env(), info, instantiate_msg)
     }
@@ -874,6 +880,9 @@ mod tests {
             },
             max_voting_period: Duration::Time(3600),
             admin_voting_threshold_percentage: 75,
+            withdrawn_staking_rewards: None,
+            withdrawn_locked: None,
+            withdrawn_unlocked: None,
         };
         let err =
             instantiate(deps.as_mut(), mock_env(), info.clone(), instantiate_msg).unwrap_err();
@@ -892,6 +901,9 @@ mod tests {
             },
             max_voting_period: Duration::Time(3600),
             admin_voting_threshold_percentage: 75,
+            withdrawn_staking_rewards: None,
+            withdrawn_locked: None,
+            withdrawn_unlocked: None,
         };
         let err =
             instantiate(deps.as_mut(), mock_env(), info.clone(), instantiate_msg).unwrap_err();
@@ -910,6 +922,9 @@ mod tests {
             },
             max_voting_period: Duration::Time(3600),
             admin_voting_threshold_percentage: 75,
+            withdrawn_staking_rewards: None,
+            withdrawn_locked: None,
+            withdrawn_unlocked: None,
         };
         let err =
             instantiate(deps.as_mut(), mock_env(), info.clone(), instantiate_msg).unwrap_err();
diff --git a/src/msg.rs b/src/msg.rs
index 81246b6..6d548ca 100644
--- a/src/msg.rs
+++ b/src/msg.rs
@@ -14,6 +14,11 @@ pub struct InstantiateMsg {
     pub tranche: Tranche,
     pub max_voting_period: Duration,
     pub admin_voting_threshold_percentage: u8,
+
+    // override cumulative state (for inheriting a contract or testing)
+    pub withdrawn_staking_rewards: Option<u128>,
+    pub withdrawn_unlocked: Option<u128>,
+    pub withdrawn_locked: Option<u128>,
 }
 
 #[cw_serde]