refactor: cleanup WalletBalance

dash-spv/src/main.rs

@@ -451,7 +451,7 @@ async fn run_client<S: dash_spv::storage::StorageManager>(
                     _ = 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<S: dash_spv::storage::StorageManager>(
                                     .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();
                         }

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

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));

key-wallet-ffi/src/types.rs

@@ -64,10 +64,10 @@ pub struct FFIBalance {
 impl From<key_wallet::WalletBalance> 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(),
         }
     }
 }

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

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);

key-wallet-manager/src/wallet_manager/mod.rs

@@ -880,7 +880,7 @@ impl<T: WalletInfoInterface> WalletManager<T> {
 
     /// 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

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();

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 {

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

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()
         );
     }

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"
         );
     }

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"
         );