diff --git a/key-wallet/src/managed_account/mod.rs b/key-wallet/src/managed_account/mod.rs index 7ebd64c41..c19e4c9cb 100644 --- a/key-wallet/src/managed_account/mod.rs +++ b/key-wallet/src/managed_account/mod.rs @@ -265,7 +265,20 @@ impl ManagedAccount { } /// Update the account balance - pub fn update_balance(&mut self, spendable: u64, unconfirmed: u64, locked: u64) { + pub fn update_balance(&mut self) { + let mut spendable = 0; + let mut unconfirmed = 0; + let mut locked = 0; + for utxo in self.utxos.values() { + let value = utxo.txout.value; + if utxo.is_locked { + locked += value; + } else if utxo.is_confirmed { + spendable += value; + } else { + unconfirmed += value; + } + } self.balance = WalletBalance::new(spendable, unconfirmed, locked); self.metadata.last_used = Some(Self::current_timestamp()); } diff --git a/key-wallet/src/transaction_checking/wallet_checker.rs b/key-wallet/src/transaction_checking/wallet_checker.rs index f23b431dc..036faadce 100644 --- a/key-wallet/src/transaction_checking/wallet_checker.rs +++ b/key-wallet/src/transaction_checking/wallet_checker.rs @@ -235,24 +235,6 @@ impl WalletTransactionChecker for ManagedWalletInfo { for input in &tx.input { account.utxos.remove(&input.previous_output); } - - // Recalculate account balance from UTXOs - let mut confirmed = 0u64; - let mut unconfirmed = 0u64; - let mut locked = 0u64; - - for utxo in account.utxos.values() { - let value = utxo.txout.value; - if utxo.is_locked { - locked += value; - } else if utxo.is_confirmed { - confirmed += value; - } else { - unconfirmed += value; - } - } - - let _ = account.update_balance(confirmed, unconfirmed, locked); } _ => { // Skip UTXO ingestion for identity/provider accounts diff --git a/key-wallet/src/wallet/balance.rs b/key-wallet/src/wallet/balance.rs index 38de7f07c..08467aa6c 100644 --- a/key-wallet/src/wallet/balance.rs +++ b/key-wallet/src/wallet/balance.rs @@ -3,6 +3,7 @@ //! This module provides a wallet balance structure containing all available balances. use core::fmt::{Display, Formatter}; +use core::ops::AddAssign; #[cfg(feature = "serde")] use serde::{Deserialize, Serialize}; @@ -62,6 +63,14 @@ impl Display for WalletBalance { } } +impl AddAssign for WalletBalance { + fn add_assign(&mut self, other: Self) { + self.spendable += other.spendable; + self.unconfirmed += other.unconfirmed; + self.locked += other.locked; + } +} + #[cfg(test)] mod tests { use super::*; @@ -91,4 +100,20 @@ mod tests { let display = balance.to_string(); assert_eq!(display, "Spendable: 1000, Unconfirmed: 500, Locked: 200, Total: 1700"); } + + #[test] + fn test_balance_add_assign() { + let mut balance = WalletBalance::new(1000, 500, 200); + let balance_add = WalletBalance::new(300, 100, 50); + // Test adding actual balances + balance += balance_add; + assert_eq!(balance.spendable(), 1300); + assert_eq!(balance.unconfirmed(), 600); + assert_eq!(balance.locked(), 250); + assert_eq!(balance.total(), 2150); + // Test adding zero balances + let balance_before = balance; + balance += WalletBalance::default(); + assert_eq!(balance_before, balance); + } } diff --git a/key-wallet/src/wallet/managed_wallet_info/wallet_info_interface.rs b/key-wallet/src/wallet/managed_wallet_info/wallet_info_interface.rs index 58e271f7d..223f575f9 100644 --- a/key-wallet/src/wallet/managed_wallet_info/wallet_info_interface.rs +++ b/key-wallet/src/wallet/managed_wallet_info/wallet_info_interface.rs @@ -3,6 +3,7 @@ //! This trait allows WalletManager to work with different wallet info implementations use super::managed_account_operations::ManagedAccountOperations; +use crate::account::ManagedAccountTrait; use crate::managed_account::managed_account_collection::ManagedAccountCollection; use crate::transaction_checking::WalletTransactionChecker; use crate::wallet::immature_transaction::{ImmatureTransaction, ImmatureTransactionCollection}; @@ -202,24 +203,12 @@ impl WalletInfoInterface for ManagedWalletInfo { } fn update_balance(&mut self) { - let mut spendable = 0u64; - let mut unconfirmed = 0u64; - let mut locked = 0u64; - - for account in self.accounts.all_accounts() { - for utxo in account.utxos.values() { - let value = utxo.txout.value; - if utxo.is_locked { - locked += value; - } else if utxo.is_confirmed { - spendable += value; - } else { - unconfirmed += value; - } - } + let mut balance = WalletBalance::default(); + for account in self.accounts.all_accounts_mut() { + account.update_balance(); + balance += *account.balance(); } - - self.balance = WalletBalance::new(spendable, unconfirmed, locked) + self.balance = balance; } fn transaction_history(&self) -> Vec<&TransactionRecord> {