fix: integrate staked tokens for governance voting power

contracts/src/goal.rs

@@ -94,6 +94,20 @@ pub fn create_goal_save(
     add_goal_to_user(env, &user, goal_id);
     increment_next_goal_id(env);
 
+    // Update user's total balance
+    let user_key = DataKey::User(user.clone());
+    if let Some(mut user_data) = env.storage().persistent().get::<DataKey, User>(&user_key) {
+        user_data.total_balance = user_data
+            .total_balance
+            .checked_add(net_initial_deposit)
+            .ok_or(SavingsError::Overflow)?;
+        user_data.savings_count = user_data
+            .savings_count
+            .checked_add(1)
+            .ok_or(SavingsError::Overflow)?;
+        env.storage().persistent().set(&user_key, &user_data);
+    }
+
     // Award deposit points
     storage::award_deposit_points(env, user.clone(), initial_deposit)?;
 
@@ -153,6 +167,16 @@ pub fn deposit_to_goal_save(
         .persistent()
         .set(&DataKey::GoalSave(goal_id), &goal_save);
 
+    // Update user's total balance
+    let user_key = DataKey::User(user.clone());
+    if let Some(mut user_data) = env.storage().persistent().get::<DataKey, User>(&user_key) {
+        user_data.total_balance = user_data
+            .total_balance
+            .checked_add(net_amount)
+            .ok_or(SavingsError::Overflow)?;
+        env.storage().persistent().set(&user_key, &user_data);
+    }
+
     if !was_completed && goal_save.is_completed {
         storage::award_goal_completion_bonus(env, user.clone())?;
     }
@@ -240,8 +264,8 @@ pub fn withdraw_completed_goal_save(
     if let Some(mut user_data) = env.storage().persistent().get::<DataKey, User>(&user_key) {
         user_data.total_balance = user_data
             .total_balance
-            .checked_add(net_amount)
-            .ok_or(SavingsError::Overflow)?;
+            .checked_sub(goal_save.current_amount)
+            .ok_or(SavingsError::Underflow)?;
         env.storage().persistent().set(&user_key, &user_data);
     }
 
@@ -335,8 +359,8 @@ pub fn break_goal_save(env: &Env, user: Address, goal_id: u64) -> Result<i128, S
     if let Some(mut user_data) = env.storage().persistent().get::<DataKey, User>(&user_key) {
         user_data.total_balance = user_data
             .total_balance
-            .checked_add(net_amount)
-            .ok_or(SavingsError::Overflow)?;
+            .checked_sub(goal_save.current_amount)
+            .ok_or(SavingsError::Underflow)?;
         env.storage().persistent().set(&user_key, &user_data);
     }
 

contracts/src/governance.rs

@@ -104,10 +104,13 @@ pub enum ProposalAction {
     UnpauseContract,
 }
 
-/// Calculates voting power for a user based on their lifetime deposited funds
+/// Calculates voting power for a user based on their current staked tokens
 pub fn get_voting_power(env: &Env, user: &Address) -> u128 {
-    let rewards = get_user_rewards(env, user.clone());
-    rewards.lifetime_deposited.max(0) as u128
+    let user_data = crate::users::get_user(env, user).unwrap_or(crate::storage_types::User {
+        total_balance: 0,
+        savings_count: 0,
+    });
+    user_data.total_balance.max(0) as u128
 }
 
 /// Creates a new governance proposal

contracts/src/governance_tests.rs

@@ -136,6 +136,20 @@ mod governance_tests {
         assert_eq!(power, 1500);
     }
 
+    #[test]
+    fn test_voting_power_decreases_after_withdraw() {
+        let (env, client, _) = setup_contract();
+        let user = Address::generate(&env);
+        env.mock_all_auths();
+
+        client.initialize_user(&user);
+        client.deposit_flexi(&user, &2000);
+        assert_eq!(client.get_voting_power(&user), 2000);
+
+        client.withdraw_flexi(&user, &500);
+        assert_eq!(client.get_voting_power(&user), 1500);
+    }
+
     #[test]
     fn test_init_voting_config() {
         let (env, client, admin) = setup_contract();

contracts/src/group.rs

@@ -452,6 +452,16 @@ pub fn contribute_to_group_save(
         env.storage().persistent().set(&plan_key, &plan);
     }
 
+    // Update user's total balance
+    let user_key = DataKey::User(user.clone());
+    if let Some(mut user_data) = env.storage().persistent().get::<DataKey, crate::storage_types::User>(&user_key) {
+        user_data.total_balance = user_data
+            .total_balance
+            .checked_add(amount)
+            .ok_or(SavingsError::Overflow)?;
+        env.storage().persistent().set(&user_key, &user_data);
+    }
+
     // Award deposit points
     crate::rewards::storage::award_deposit_points(env, user.clone(), amount)?;
 
@@ -631,6 +641,16 @@ pub fn break_group_save(env: &Env, user: Address, group_id: u64) -> Result<(), S
     // Save updated group
     env.storage().persistent().set(&group_key, &group);
 
+    // Update user's total balance
+    let user_key = DataKey::User(user.clone());
+    if let Some(mut user_data) = env.storage().persistent().get::<DataKey, crate::storage_types::User>(&user_key) {
+        user_data.total_balance = user_data
+            .total_balance
+            .checked_sub(user_contribution)
+            .ok_or(SavingsError::Underflow)?;
+        env.storage().persistent().set(&user_key, &user_data);
+    }
+
     // Remove user's contribution entry
     env.storage().persistent().remove(&contribution_key);
 

contracts/src/staking/storage.rs

@@ -177,8 +177,7 @@ pub fn update_rewards(env: &Env) -> Result<(), SavingsError> {
 
     if total_staked > 0 {
         // For first stake (last_update == 0), use current time as reference
-        let effective_last_update = if last_update == 0 { now } else { last_update };
-
+
         let time_elapsed = now
             .checked_sub(effective_last_update)
             .ok_or(SavingsError::Underflow)?;

contracts/src/staking_tests.rs

@@ -14,11 +14,9 @@ fn setup_env_with_staking(config: StakingConfig) -> (Env, NesteraContractClient<
     let admin = Address::generate(&env);
     let admin_pk = BytesN::from_array(&env, &[9u8; 32]);
 
-    env.mock_all_auths();
-
+
     // Set initial ledger timestamp to non-zero value
-    env.ledger().with_mut(|li| li.timestamp = 1000);
-
+
     client.initialize(&admin, &admin_pk);
     assert!(client.try_init_staking_config(&admin, &config).is_ok());