diff --git a/contracts/src/goal.rs b/contracts/src/goal.rs index 12ce1a4da..5ae272286 100644 --- a/contracts/src/goal.rs +++ b/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::(&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::(&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::(&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(&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); } diff --git a/contracts/src/governance.rs b/contracts/src/governance.rs index 9ba62c97d..b458f4201 100644 --- a/contracts/src/governance.rs +++ b/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 diff --git a/contracts/src/governance_tests.rs b/contracts/src/governance_tests.rs index 9e1ef2671..2383ffe2f 100644 --- a/contracts/src/governance_tests.rs +++ b/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(); diff --git a/contracts/src/group.rs b/contracts/src/group.rs index e199de594..324d501f0 100644 --- a/contracts/src/group.rs +++ b/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::(&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::(&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); diff --git a/contracts/src/staking/storage.rs b/contracts/src/staking/storage.rs index 54bd30361..c7bbdc225 100644 --- a/contracts/src/staking/storage.rs +++ b/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)?; diff --git a/contracts/src/staking_tests.rs b/contracts/src/staking_tests.rs index 11ce15250..1624d5cee 100644 --- a/contracts/src/staking_tests.rs +++ b/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());