From 88e29657e5268d0a3a45194c252192bea35c3c9c Mon Sep 17 00:00:00 2001
From: Noah Saso <noahsaso@gmail.com>
Date: Tue, 21 May 2024 22:45:33 -0400
Subject: [PATCH] increase time before querying total staked, and fixed height
 in execute fn

---
 .../voting/dao-voting-cosmos-staked/src/contract.rs    |  5 +++--
 .../src/tests/test_tube/integration_tests.rs           | 10 +++++-----
 2 files changed, 8 insertions(+), 7 deletions(-)

diff --git a/contracts/voting/dao-voting-cosmos-staked/src/contract.rs b/contracts/voting/dao-voting-cosmos-staked/src/contract.rs
index 2cc164ed8..e38c8f733 100644
--- a/contracts/voting/dao-voting-cosmos-staked/src/contract.rs
+++ b/contracts/voting/dao-voting-cosmos-staked/src/contract.rs
@@ -54,12 +54,13 @@ pub fn execute_update_total_staked(
     amount: Uint128,
     height: Option<u64>,
 ) -> Result<Response, ContractError> {
-    STAKED_TOTAL.save(deps.storage, &amount, env.block.height)?;
+    let height = height.unwrap_or(env.block.height);
+    STAKED_TOTAL.save(deps.storage, &amount, height)?;
 
     Ok(Response::new()
         .add_attribute("action", "update_total_staked")
         .add_attribute("amount", amount)
-        .add_attribute("height", height.unwrap_or(env.block.height).to_string()))
+        .add_attribute("height", height.to_string()))
 }
 
 #[cfg_attr(not(feature = "library"), entry_point)]
diff --git a/contracts/voting/dao-voting-cosmos-staked/src/tests/test_tube/integration_tests.rs b/contracts/voting/dao-voting-cosmos-staked/src/tests/test_tube/integration_tests.rs
index c1b3ef53e..1ea6f3006 100644
--- a/contracts/voting/dao-voting-cosmos-staked/src/tests/test_tube/integration_tests.rs
+++ b/contracts/voting/dao-voting-cosmos-staked/src/tests/test_tube/integration_tests.rs
@@ -82,8 +82,8 @@ fn test_staked_voting_power_and_update() {
                     msgs: vec![CosmosMsg::Stargate {
                         type_url: "/cosmos.authz.v1beta1.MsgGrant".to_string(),
                         value: MsgGrant {
-                            granter: dao.contract_addr.to_string(),
-                            grantee: bot.address().to_string(),
+                            granter: dao.contract_addr.clone(),
+                            grantee: bot.address(),
                             grant: Some(Grant {
                                 authorization: Some(
                                     ContractExecutionAuthorization {
@@ -119,9 +119,6 @@ fn test_staked_voting_power_and_update() {
             staker,
         )
         .unwrap();
-
-    app.increase_time(10);
-
     proposal_single
         .execute(
             &dao_proposal_single::msg::ExecuteMsg::Execute { proposal_id: 1 },
@@ -152,6 +149,9 @@ fn test_staked_voting_power_and_update() {
         )
         .unwrap();
 
+    // Move chain forward so we can update total staked on a new block
+    app.increase_time(100);
+
     // Query total power
     let total_power = vp_contract.query_tp(None).unwrap();
     assert_eq!(total_power.power, Uint128::new(100));