diff --git a/contracts/stake/src/contract.rs b/contracts/stake/src/contract.rs
index f75cb355..e34f387f 100644
--- a/contracts/stake/src/contract.rs
+++ b/contracts/stake/src/contract.rs
@@ -375,10 +375,11 @@ impl Staking {
         admin.require_auth();
         env.deployer().update_current_contract_wasm(new_wasm_hash);
 
-        env.storage()
-            .persistent()
-            .set(&DistributionDataKey::TotalStakedHistory, &map![&env, (0, 0)]);
-
+        let current_timestamp = env.ledger().timestamp();
+        let total_staked_amount = get_total_staked_counter(&env);
+        let mut total_staked_history = map![&env];
+        total_staked_history.set(current_timestamp, total_staked_amount as u128);
+        save_total_staked_history(&env, total_staked_history);
     }
 }
 
diff --git a/contracts/stake/src/tests/setup.rs b/contracts/stake/src/tests/setup.rs
index aa720002..4d9463af 100644
--- a/contracts/stake/src/tests/setup.rs
+++ b/contracts/stake/src/tests/setup.rs
@@ -117,6 +117,7 @@ fn upgrade_stake_contract() {
 
     let new_stake_wasm = install_stake_latest_wasm(&env);
     stake_v_1_0_0_client.update(&new_stake_wasm);
+    stake_v_1_0_0_client.update(&new_stake_wasm);
 
     let upgraded_stake_client = stake_latest::Client::new(&env, &stake_addr);
 
@@ -134,10 +135,6 @@ fn upgrade_stake_contract() {
     );
 
     upgraded_stake_client.create_distribution_flow(&owner, &token_client.address);
-    token_client.mint(&admin, &1_000);
-    dbg!("Query now:");
-    dbg!(upgraded_stake_client.query_total_staked_history());
-    // dbg!("here?");
-    // dbg!(upgraded_stake_client.query_total_staked_history());
-    // upgraded_stake_client.distribute_rewards(&owner, &1_000, &token_client.address);
+    token_client.mint(&owner, &1_000);
+    upgraded_stake_client.distribute_rewards(&owner, &1_000, &token_client.address);
 }