From 2860407657b6b156f204c9387374ce1ff8e1a096 Mon Sep 17 00:00:00 2001 From: xdustinface Date: Sat, 3 Jan 2026 03:43:33 +0100 Subject: [PATCH] refactor: cleanup `WalletBalance` - some renaming of the fields - make the fields private - drop overflow checks - implement, test and use `core::fmt::Display` - simplify tests - drop unused/test-only code --- dash-spv/src/main.rs | 14 +- key-wallet-ffi/src/managed_account.rs | 6 +- key-wallet-ffi/src/managed_wallet.rs | 15 +- key-wallet-ffi/src/types.rs | 8 +- key-wallet-ffi/src/wallet_manager.rs | 4 +- .../examples/wallet_creation.rs | 2 +- key-wallet-manager/src/wallet_manager/mod.rs | 2 +- key-wallet-manager/tests/integration_test.rs | 4 +- key-wallet/src/lib.rs | 5 +- key-wallet/src/managed_account/mod.rs | 10 +- .../transaction_router/tests/coinbase.rs | 7 +- .../transaction_router/tests/routing.rs | 3 +- .../transaction_checking/wallet_checker.rs | 2 +- key-wallet/src/wallet/balance.rs | 394 ++---------------- .../src/wallet/managed_wallet_info/mod.rs | 24 -- .../wallet_info_interface.rs | 8 +- key-wallet/src/wallet/mod.rs | 2 +- 17 files changed, 83 insertions(+), 427 deletions(-) diff --git a/dash-spv/src/main.rs b/dash-spv/src/main.rs index 51c8207ba..8f98a3c46 100644 --- a/dash-spv/src/main.rs +++ b/dash-spv/src/main.rs @@ -451,7 +451,7 @@ async fn run_client( _ = tokio::time::sleep(snapshot_interval) => { // Log snapshot if interval has elapsed if last_snapshot.elapsed() >= snapshot_interval { - let (tx_count, confirmed, unconfirmed, locked, total) = { + let (tx_count, wallet_balance) = { let mgr = wallet_for_logger.read().await; // Count wallet-affecting transactions from wallet transaction history @@ -461,20 +461,16 @@ async fn run_client( .unwrap_or(0); // Read wallet balance from the managed wallet info - let wb = mgr.get_wallet_balance(&wallet_id_for_logger).ok(); - let (c, u, l, t) = wb.map(|b| (b.confirmed, b.unconfirmed, b.locked, b.total)).unwrap_or((0, 0, 0, 0)); + let wallet_balance = mgr.get_wallet_balance(&wallet_id_for_logger).unwrap_or_default(); - (tx_count, c, u, l, t) + (tx_count, wallet_balance) }; tracing::info!( - "Wallet tx summary: tx_count={} (blocks={} + mempool={}), balances: confirmed={} unconfirmed={} locked={} total={}", + "Wallet tx summary: tx_count={} (blocks={} + mempool={}), balances: {}", tx_count, total_detected_block_txs, total_detected_mempool_txs, - confirmed, - unconfirmed, - locked, - total, + wallet_balance, ); last_snapshot = std::time::Instant::now(); } diff --git a/key-wallet-ffi/src/managed_account.rs b/key-wallet-ffi/src/managed_account.rs index 92a8e28fc..10d7afee3 100644 --- a/key-wallet-ffi/src/managed_account.rs +++ b/key-wallet-ffi/src/managed_account.rs @@ -513,10 +513,10 @@ pub unsafe extern "C" fn managed_account_get_balance( let balance = &account.inner().balance; *balance_out = crate::types::FFIBalance { - confirmed: balance.confirmed, - unconfirmed: balance.unconfirmed, + confirmed: balance.spendable(), + unconfirmed: balance.unconfirmed(), immature: 0, // WalletBalance doesn't have immature field - total: balance.total, + total: balance.total(), }; true diff --git a/key-wallet-ffi/src/managed_wallet.rs b/key-wallet-ffi/src/managed_wallet.rs index f807349b5..81a3f111a 100644 --- a/key-wallet-ffi/src/managed_wallet.rs +++ b/key-wallet-ffi/src/managed_wallet.rs @@ -576,10 +576,10 @@ pub unsafe extern "C" fn managed_wallet_get_balance( let balance = &managed_wallet.inner().balance; unsafe { - *confirmed_out = balance.confirmed; - *unconfirmed_out = balance.unconfirmed; - *locked_out = balance.locked; - *total_out = balance.total; + *confirmed_out = balance.spendable(); + *unconfirmed_out = balance.unconfirmed(); + *locked_out = balance.locked(); + *total_out = balance.total(); } FFIError::set_success(error); @@ -1042,12 +1042,7 @@ mod tests { let mut managed_info = ManagedWalletInfo::from_wallet(wallet_arc); // Set some test balance values - managed_info.balance = WalletBalance { - confirmed: 1000000, - unconfirmed: 50000, - locked: 25000, - total: 1075000, - }; + managed_info.balance = WalletBalance::new(1000000, 50000, 25000); let ffi_managed = FFIManagedWalletInfo::new(managed_info); let ffi_managed_ptr = Box::into_raw(Box::new(ffi_managed)); diff --git a/key-wallet-ffi/src/types.rs b/key-wallet-ffi/src/types.rs index 90f91880c..f8c4217f2 100644 --- a/key-wallet-ffi/src/types.rs +++ b/key-wallet-ffi/src/types.rs @@ -64,10 +64,10 @@ pub struct FFIBalance { impl From for FFIBalance { fn from(balance: key_wallet::WalletBalance) -> Self { FFIBalance { - confirmed: balance.confirmed, - unconfirmed: balance.unconfirmed, - immature: balance.locked, // Map locked to immature for now - total: balance.total, + confirmed: balance.spendable(), + unconfirmed: balance.unconfirmed(), + immature: balance.locked(), // Map locked to immature for now + total: balance.total(), } } } diff --git a/key-wallet-ffi/src/wallet_manager.rs b/key-wallet-ffi/src/wallet_manager.rs index 6bdec2876..6ee383ec5 100644 --- a/key-wallet-ffi/src/wallet_manager.rs +++ b/key-wallet-ffi/src/wallet_manager.rs @@ -703,8 +703,8 @@ pub unsafe extern "C" fn wallet_manager_get_wallet_balance( match result { Ok(balance) => { unsafe { - *confirmed_out = balance.confirmed; - *unconfirmed_out = balance.unconfirmed; + *confirmed_out = balance.spendable(); + *unconfirmed_out = balance.unconfirmed(); } FFIError::set_success(error); true diff --git a/key-wallet-manager/examples/wallet_creation.rs b/key-wallet-manager/examples/wallet_creation.rs index 0b13fb4aa..571ac3232 100644 --- a/key-wallet-manager/examples/wallet_creation.rs +++ b/key-wallet-manager/examples/wallet_creation.rs @@ -127,7 +127,7 @@ fn main() { for (i, wallet_id) in [wallet_id, wallet_id2].iter().enumerate() { match manager.get_wallet_balance(wallet_id) { Ok(balance) => { - println!(" Wallet {}: {} satoshis", i + 1, balance.total); + println!(" Wallet {}: {} satoshis", i + 1, balance.total()); } Err(e) => { println!(" Wallet {}: Error - {:?}", i + 1, e); diff --git a/key-wallet-manager/src/wallet_manager/mod.rs b/key-wallet-manager/src/wallet_manager/mod.rs index d52b8c368..0e4598810 100644 --- a/key-wallet-manager/src/wallet_manager/mod.rs +++ b/key-wallet-manager/src/wallet_manager/mod.rs @@ -880,7 +880,7 @@ impl WalletManager { /// Get total balance across all wallets and networks pub fn get_total_balance(&self) -> u64 { - self.wallet_infos.values().map(|info| info.balance().total).sum() + self.wallet_infos.values().map(|info| info.balance().total()).sum() } /// Get balance for a specific wallet diff --git a/key-wallet-manager/tests/integration_test.rs b/key-wallet-manager/tests/integration_test.rs index 4ede4e4fd..31e31ace8 100644 --- a/key-wallet-manager/tests/integration_test.rs +++ b/key-wallet-manager/tests/integration_test.rs @@ -126,7 +126,7 @@ fn test_utxo_management() { let balance = manager.get_wallet_balance(&wallet_id); assert!(balance.is_ok()); - assert_eq!(balance.unwrap().total, 0); + assert_eq!(balance.unwrap().total(), 0); } #[test] @@ -145,7 +145,7 @@ fn test_balance_calculation() { // Check wallet balance (should be 0 initially) let balance = manager.get_wallet_balance(&wallet_id); assert!(balance.is_ok()); - assert_eq!(balance.unwrap().total, 0); + assert_eq!(balance.unwrap().total(), 0); // Check global balance let total = manager.get_total_balance(); diff --git a/key-wallet/src/lib.rs b/key-wallet/src/lib.rs index 08e957229..2ecf1eb36 100644 --- a/key-wallet/src/lib.rs +++ b/key-wallet/src/lib.rs @@ -65,10 +65,7 @@ pub use managed_account::managed_account_type::ManagedAccountType; pub use mnemonic::Mnemonic; pub use seed::Seed; pub use utxo::{Utxo, UtxoSet}; -pub use wallet::{ - balance::{BalanceError, WalletBalance}, - Wallet, -}; +pub use wallet::{balance::WalletBalance, Wallet}; /// Re-export commonly used types pub mod prelude { diff --git a/key-wallet/src/managed_account/mod.rs b/key-wallet/src/managed_account/mod.rs index 94fea9d6f..7ebd64c41 100644 --- a/key-wallet/src/managed_account/mod.rs +++ b/key-wallet/src/managed_account/mod.rs @@ -265,15 +265,9 @@ impl ManagedAccount { } /// Update the account balance - pub fn update_balance( - &mut self, - confirmed: u64, - unconfirmed: u64, - locked: u64, - ) -> Result<(), crate::wallet::balance::BalanceError> { - self.balance.update(confirmed, unconfirmed, locked)?; + pub fn update_balance(&mut self, spendable: u64, unconfirmed: u64, locked: u64) { + self.balance = WalletBalance::new(spendable, unconfirmed, locked); self.metadata.last_used = Some(Self::current_timestamp()); - Ok(()) } /// Get all addresses from all pools diff --git a/key-wallet/src/transaction_checking/transaction_router/tests/coinbase.rs b/key-wallet/src/transaction_checking/transaction_router/tests/coinbase.rs index eb937c489..a969cfae2 100644 --- a/key-wallet/src/transaction_checking/transaction_router/tests/coinbase.rs +++ b/key-wallet/src/transaction_checking/transaction_router/tests/coinbase.rs @@ -230,7 +230,7 @@ async fn test_update_state_flag_behavior() { let address = managed_account .next_receive_address(Some(&xpub), true) .expect("Failed to generate receive address"); - let balance = managed_account.balance.confirmed; + let balance = managed_account.balance.spendable(); let tx_count = managed_account.transactions.len(); (address, balance, tx_count) }; @@ -261,7 +261,8 @@ async fn test_update_state_flag_behavior() { .first_bip44_managed_account_mut() .expect("Failed to get first BIP44 managed account"); assert_eq!( - managed_account.balance.confirmed, initial_balance, + managed_account.balance.spendable(), + initial_balance, "Balance should not change when update_state=false" ); assert_eq!( @@ -296,7 +297,7 @@ async fn test_update_state_flag_behavior() { .expect("Failed to get first BIP44 managed account"); println!( "After update_state=true: balance={}, tx_count={}", - managed_account.balance.confirmed, + managed_account.balance.spendable(), managed_account.transactions.len() ); } diff --git a/key-wallet/src/transaction_checking/transaction_router/tests/routing.rs b/key-wallet/src/transaction_checking/transaction_router/tests/routing.rs index 8379735b8..34f9b54f2 100644 --- a/key-wallet/src/transaction_checking/transaction_router/tests/routing.rs +++ b/key-wallet/src/transaction_checking/transaction_router/tests/routing.rs @@ -169,7 +169,8 @@ async fn test_transaction_routing_to_bip32_account() { .first_bip32_managed_account_mut() .expect("Failed to get first BIP32 managed account"); assert_eq!( - managed_account.balance.confirmed, 0, + managed_account.balance.spendable(), + 0, "Balance should not be updated when update_state is false" ); } diff --git a/key-wallet/src/transaction_checking/wallet_checker.rs b/key-wallet/src/transaction_checking/wallet_checker.rs index 8ce9d6b23..f23b431dc 100644 --- a/key-wallet/src/transaction_checking/wallet_checker.rs +++ b/key-wallet/src/transaction_checking/wallet_checker.rs @@ -813,7 +813,7 @@ mod tests { // Verify balance is still zero assert_eq!( - managed_wallet.balance().total, + managed_wallet.balance().total(), 0, "Balance should be zero while coinbase is immature" ); diff --git a/key-wallet/src/wallet/balance.rs b/key-wallet/src/wallet/balance.rs index 9d9593191..38de7f07c 100644 --- a/key-wallet/src/wallet/balance.rs +++ b/key-wallet/src/wallet/balance.rs @@ -1,9 +1,8 @@ -//! Wallet balance management +//! Wallet balance //! -//! This module provides wallet balance tracking and state transition functionality -//! for managing confirmed, unconfirmed, and locked balances. +//! This module provides a wallet balance structure containing all available balances. -use alloc::string::String; +use core::fmt::{Display, Formatter}; #[cfg(feature = "serde")] use serde::{Deserialize, Serialize}; @@ -11,386 +10,85 @@ use serde::{Deserialize, Serialize}; #[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] pub struct WalletBalance { - /// Confirmed balance (UTXOs with confirmations) - pub confirmed: u64, + /// Confirmed and mature balance (UTXOs with enough confirmations to be spendable) + spendable: u64, /// Unconfirmed balance (UTXOs without confirmations) - pub unconfirmed: u64, + unconfirmed: u64, /// Locked balance (UTXOs reserved for specific purposes like CoinJoin) - pub locked: u64, - /// Total balance (sum of all balances) - pub total: u64, + locked: u64, } impl WalletBalance { /// Create a new wallet balance - pub fn new(confirmed: u64, unconfirmed: u64, locked: u64) -> Result { - let total = confirmed - .checked_add(unconfirmed) - .and_then(|sum| sum.checked_add(locked)) - .ok_or(BalanceError::Overflow)?; - - Ok(Self { - confirmed, + pub fn new(spendable: u64, unconfirmed: u64, locked: u64) -> Self { + Self { + spendable, unconfirmed, locked, - total, - }) - } - - /// Create an empty balance - pub fn zero() -> Self { - Self::default() + } } - /// Get spendable balance (confirmed only, excluding locked) + /// Get the spendable balance. pub fn spendable(&self) -> u64 { - self.confirmed + self.spendable } - /// Get pending balance (unconfirmed) - pub fn pending(&self) -> u64 { + /// Get the unconfirmed balance. + pub fn unconfirmed(&self) -> u64 { self.unconfirmed } - /// Get available balance (confirmed + unconfirmed, excluding locked) - pub fn available(&self) -> u64 { - self.confirmed + self.unconfirmed - } - - /// Mature locked balance by moving an amount from locked to confirmed - /// This happens when locked funds (e.g., from CoinJoin) become available - pub fn mature(&mut self, amount: u64) -> Result<(), BalanceError> { - if amount > self.locked { - return Err(BalanceError::InsufficientLockedBalance { - requested: amount, - available: self.locked, - }); - } - - self.locked = self.locked.checked_sub(amount).ok_or(BalanceError::Underflow)?; - self.confirmed = self.confirmed.checked_add(amount).ok_or(BalanceError::Overflow)?; - // Total remains the same - Ok(()) - } - - /// Confirm unconfirmed balance by moving an amount from unconfirmed to confirmed - /// This happens when transactions get confirmed in blocks - pub fn confirm(&mut self, amount: u64) -> Result<(), BalanceError> { - if amount > self.unconfirmed { - return Err(BalanceError::InsufficientUnconfirmedBalance { - requested: amount, - available: self.unconfirmed, - }); - } - - self.unconfirmed = self.unconfirmed.checked_sub(amount).ok_or(BalanceError::Underflow)?; - self.confirmed = self.confirmed.checked_add(amount).ok_or(BalanceError::Overflow)?; - // Total remains the same - Ok(()) - } - - /// Lock confirmed balance by moving an amount from confirmed to locked - /// This happens when funds are reserved for specific purposes - pub fn lock(&mut self, amount: u64) -> Result<(), BalanceError> { - if amount > self.confirmed { - return Err(BalanceError::InsufficientConfirmedBalance { - requested: amount, - available: self.confirmed, - }); - } - - self.confirmed = self.confirmed.checked_sub(amount).ok_or(BalanceError::Underflow)?; - self.locked = self.locked.checked_add(amount).ok_or(BalanceError::Overflow)?; - // Total remains the same - Ok(()) - } - - /// Add incoming unconfirmed balance - pub fn add_unconfirmed(&mut self, amount: u64) -> Result<(), BalanceError> { - self.unconfirmed = self.unconfirmed.checked_add(amount).ok_or(BalanceError::Overflow)?; - self.total = self.total.checked_add(amount).ok_or(BalanceError::Overflow)?; - Ok(()) + /// Get the locked balance. + pub fn locked(&self) -> u64 { + self.locked } - /// Add incoming confirmed balance - pub fn add_confirmed(&mut self, amount: u64) -> Result<(), BalanceError> { - self.confirmed = self.confirmed.checked_add(amount).ok_or(BalanceError::Overflow)?; - self.total = self.total.checked_add(amount).ok_or(BalanceError::Overflow)?; - Ok(()) - } - - /// Remove spent confirmed balance - pub fn spend_confirmed(&mut self, amount: u64) -> Result<(), BalanceError> { - if amount > self.confirmed { - return Err(BalanceError::InsufficientConfirmedBalance { - requested: amount, - available: self.confirmed, - }); - } - - self.confirmed = self.confirmed.checked_sub(amount).ok_or(BalanceError::Underflow)?; - self.total = self.total.checked_sub(amount).ok_or(BalanceError::Underflow)?; - Ok(()) - } - - /// Remove spent unconfirmed balance (e.g., double-spend or replacement) - pub fn remove_unconfirmed(&mut self, amount: u64) -> Result<(), BalanceError> { - if amount > self.unconfirmed { - return Err(BalanceError::InsufficientUnconfirmedBalance { - requested: amount, - available: self.unconfirmed, - }); - } - - self.unconfirmed = self.unconfirmed.checked_sub(amount).ok_or(BalanceError::Underflow)?; - self.total = self.total.checked_sub(amount).ok_or(BalanceError::Underflow)?; - Ok(()) - } - - /// Update all balance components at once - pub fn update( - &mut self, - confirmed: u64, - unconfirmed: u64, - locked: u64, - ) -> Result<(), BalanceError> { - let total = confirmed - .checked_add(unconfirmed) - .and_then(|sum| sum.checked_add(locked)) - .ok_or(BalanceError::Overflow)?; - - self.confirmed = confirmed; - self.unconfirmed = unconfirmed; - self.locked = locked; - self.total = total; - Ok(()) - } - - /// Check if balance is empty - pub fn is_empty(&self) -> bool { - self.total == 0 - } - - /// Format balance as a human-readable string - pub fn format_display(&self) -> String { - use alloc::format; - format!( - "Confirmed: {}, Unconfirmed: {}, Locked: {}, Total: {}", - self.confirmed, self.unconfirmed, self.locked, self.total - ) + /// Get the total balance. + pub fn total(&self) -> u64 { + self.spendable + self.unconfirmed + self.locked } } -/// Balance operation errors -#[derive(Debug, Clone, PartialEq, Eq)] -pub enum BalanceError { - /// Insufficient confirmed balance for operation - InsufficientConfirmedBalance { - requested: u64, - available: u64, - }, - /// Insufficient unconfirmed balance for operation - InsufficientUnconfirmedBalance { - requested: u64, - available: u64, - }, - /// Insufficient locked balance for operation - InsufficientLockedBalance { - requested: u64, - available: u64, - }, - /// Arithmetic overflow occurred - Overflow, - /// Arithmetic underflow occurred - Underflow, -} - -impl core::fmt::Display for BalanceError { - fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { - match self { - BalanceError::InsufficientConfirmedBalance { - requested, - available, - } => { - write!( - f, - "Insufficient confirmed balance: requested {} but only {} available", - requested, available - ) - } - BalanceError::InsufficientUnconfirmedBalance { - requested, - available, - } => { - write!( - f, - "Insufficient unconfirmed balance: requested {} but only {} available", - requested, available - ) - } - BalanceError::InsufficientLockedBalance { - requested, - available, - } => { - write!( - f, - "Insufficient locked balance: requested {} but only {} available", - requested, available - ) - } - BalanceError::Overflow => { - write!(f, "Arithmetic overflow in balance calculation") - } - BalanceError::Underflow => { - write!(f, "Arithmetic underflow in balance calculation") - } - } +impl Display for WalletBalance { + fn fmt(&self, f: &mut Formatter<'_>) -> core::fmt::Result { + write!( + f, + "Spendable: {}, Unconfirmed: {}, Locked: {}, Total: {}", + self.spendable, + self.unconfirmed, + self.locked, + self.total() + ) } } -#[cfg(feature = "std")] -impl std::error::Error for BalanceError {} - #[cfg(test)] mod tests { use super::*; #[test] - fn test_balance_creation() { - let balance = WalletBalance::new(1000, 500, 200).unwrap(); - assert_eq!(balance.confirmed, 1000); - assert_eq!(balance.unconfirmed, 500); - assert_eq!(balance.locked, 200); - assert_eq!(balance.total, 1700); - } - - #[test] - fn test_balance_creation_overflow() { - let result = WalletBalance::new(u64::MAX, 1, 0); - assert_eq!(result, Err(BalanceError::Overflow)); - } - - #[test] - fn test_balance_mature() { - let mut balance = WalletBalance::new(1000, 500, 200).unwrap(); - - // Mature 100 from locked to confirmed - assert!(balance.mature(100).is_ok()); - assert_eq!(balance.confirmed, 1100); - assert_eq!(balance.locked, 100); - assert_eq!(balance.total, 1700); // Total unchanged - - // Try to mature more than available - assert!(balance.mature(200).is_err()); - } - - #[test] - fn test_balance_confirm() { - let mut balance = WalletBalance::new(1000, 500, 200).unwrap(); - - // Confirm 300 from unconfirmed to confirmed - assert!(balance.confirm(300).is_ok()); - assert_eq!(balance.confirmed, 1300); - assert_eq!(balance.unconfirmed, 200); - assert_eq!(balance.total, 1700); // Total unchanged - - // Try to confirm more than available - assert!(balance.confirm(300).is_err()); - } - - #[test] - fn test_balance_lock() { - let mut balance = WalletBalance::new(1000, 500, 200).unwrap(); - - // Lock 400 from confirmed - assert!(balance.lock(400).is_ok()); - assert_eq!(balance.confirmed, 600); - assert_eq!(balance.locked, 600); - assert_eq!(balance.total, 1700); // Total unchanged - - // Try to lock more than available - assert!(balance.lock(700).is_err()); - } - - #[test] - fn test_balance_spend() { - let mut balance = WalletBalance::new(1000, 500, 200).unwrap(); - - // Spend 400 confirmed - assert!(balance.spend_confirmed(400).is_ok()); - assert_eq!(balance.confirmed, 600); - assert_eq!(balance.total, 1300); // Total reduced - - // Try to spend more than available - assert!(balance.spend_confirmed(700).is_err()); - } - - #[test] - fn test_balance_add_remove() { - let mut balance = WalletBalance::new(1000, 0, 0).unwrap(); - - // Add unconfirmed - assert!(balance.add_unconfirmed(500).is_ok()); - assert_eq!(balance.unconfirmed, 500); - assert_eq!(balance.total, 1500); - - // Add confirmed - assert!(balance.add_confirmed(300).is_ok()); - assert_eq!(balance.confirmed, 1300); - assert_eq!(balance.total, 1800); - - // Remove unconfirmed - assert!(balance.remove_unconfirmed(200).is_ok()); - assert_eq!(balance.unconfirmed, 300); - assert_eq!(balance.total, 1600); - } - - #[test] - fn test_balance_helpers() { - let balance = WalletBalance::new(1000, 500, 200).unwrap(); - + fn test_balance_creation_and_getters() { + let balance = WalletBalance::new(1000, 500, 200); assert_eq!(balance.spendable(), 1000); - assert_eq!(balance.pending(), 500); - assert_eq!(balance.available(), 1500); - assert!(!balance.is_empty()); - - let empty = WalletBalance::zero(); - assert!(empty.is_empty()); + assert_eq!(balance.unconfirmed(), 500); + assert_eq!(balance.locked(), 200); + assert_eq!(balance.total(), 1700); } #[test] - fn test_balance_update() { - let mut balance = WalletBalance::new(1000, 500, 200).unwrap(); - - assert!(balance.update(2000, 1000, 500).is_ok()); - assert_eq!(balance.confirmed, 2000); - assert_eq!(balance.unconfirmed, 1000); - assert_eq!(balance.locked, 500); - assert_eq!(balance.total, 3500); + #[should_panic(expected = "attempt to add with overflow")] + fn test_balance_overflow() { + let balance = WalletBalance::new(u64::MAX, u64::MAX, u64::MAX); + balance.total(); } #[test] - fn test_overflow_protection() { - let mut balance = WalletBalance::new(u64::MAX - 100, 0, 0).unwrap(); + fn test_balance_display() { + let zero = WalletBalance::default(); + assert_eq!(zero.to_string(), "Spendable: 0, Unconfirmed: 0, Locked: 0, Total: 0"); - // Test overflow in add_confirmed - assert_eq!(balance.add_confirmed(200), Err(BalanceError::Overflow)); - - // Test overflow in confirm - balance.unconfirmed = 200; - balance.confirmed = u64::MAX - 100; - assert_eq!(balance.confirm(200), Err(BalanceError::Overflow)); - } - - #[test] - fn test_balance_error_display() { - let err = BalanceError::InsufficientConfirmedBalance { - requested: 1000, - available: 500, - }; - let err_str = err.to_string(); - assert!(err_str.contains("Insufficient confirmed balance")); - assert!(err_str.contains("1000")); - assert!(err_str.contains("500")); + let balance = WalletBalance::new(1000, 500, 200); + let display = balance.to_string(); + assert_eq!(display, "Spendable: 1000, Unconfirmed: 500, Locked: 200, Total: 1700"); } } diff --git a/key-wallet/src/wallet/managed_wallet_info/mod.rs b/key-wallet/src/wallet/managed_wallet_info/mod.rs index 7dbc41def..0dcf482e5 100644 --- a/key-wallet/src/wallet/managed_wallet_info/mod.rs +++ b/key-wallet/src/wallet/managed_wallet_info/mod.rs @@ -120,30 +120,6 @@ impl ManagedWalletInfo { pub fn increment_transactions(&mut self) { self.metadata.total_transactions += 1; } - - /// Get total wallet balance by recalculating from all accounts (for verification) - pub fn calculate_balance(&self) -> WalletBalance { - let mut confirmed = 0u64; - let mut unconfirmed = 0u64; - let mut locked = 0u64; - - // Sum balances from all accounts across all networks - 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 { - confirmed += value; - } else { - unconfirmed += value; - } - } - } - - WalletBalance::new(confirmed, unconfirmed, locked) - .unwrap_or_else(|_| WalletBalance::default()) - } } /// Re-export types from account module for convenience 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 556891168..58e271f7d 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 @@ -202,7 +202,7 @@ impl WalletInfoInterface for ManagedWalletInfo { } fn update_balance(&mut self) { - let mut confirmed = 0u64; + let mut spendable = 0u64; let mut unconfirmed = 0u64; let mut locked = 0u64; @@ -212,16 +212,14 @@ impl WalletInfoInterface for ManagedWalletInfo { if utxo.is_locked { locked += value; } else if utxo.is_confirmed { - confirmed += value; + spendable += value; } else { unconfirmed += value; } } } - // Update balance, ignoring overflow errors as we're recalculating from scratch - self.balance = WalletBalance::new(confirmed, unconfirmed, locked) - .unwrap_or_else(|_| WalletBalance::default()); + self.balance = WalletBalance::new(spendable, unconfirmed, locked) } fn transaction_history(&self) -> Vec<&TransactionRecord> { diff --git a/key-wallet/src/wallet/mod.rs b/key-wallet/src/wallet/mod.rs index 3623f0821..858635553 100644 --- a/key-wallet/src/wallet/mod.rs +++ b/key-wallet/src/wallet/mod.rs @@ -16,7 +16,7 @@ pub mod metadata; pub mod root_extended_keys; pub mod stats; -pub use self::balance::{BalanceError, WalletBalance}; +pub use self::balance::WalletBalance; pub use self::managed_wallet_info::ManagedWalletInfo; use self::root_extended_keys::{RootExtendedPrivKey, RootExtendedPubKey}; use crate::account::account_collection::AccountCollection;