ledger: add_address returns a new Ledger with the new address pushed …

…into it.
chainwayxyz · Jun 7, 2024 · 397ac4e · 397ac4e
1 parent 1527884
commit 397ac4e
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 7 deletions.
diff --git a/src/client/rpc_api.rs b/src/client/rpc_api.rs
@@ -145,6 +145,12 @@ impl RpcApi for Client {
             .as_unchecked()
             .to_owned();
 
+        self.ledger.set(
+            self.ledger
+                .take()
+                .add_address(address.clone().assume_checked()),
+        );
+
         Ok(address)
     }
 
@@ -283,7 +289,7 @@ mod tests {
     }
 
     #[test]
-    fn generate_new_address() {
+    fn get_new_address() {
         let rpc = Client::new("", bitcoincore_rpc::Auth::None).unwrap();
 
         let address = rpc.get_new_address(None, None).unwrap();
@@ -292,17 +298,24 @@ mod tests {
         assert!(!address.is_valid_for_network(Network::Testnet));
         assert!(!address.is_valid_for_network(Network::Signet));
         assert!(!address.is_valid_for_network(Network::Bitcoin));
+        unsafe { assert_eq!(*(*rpc.ledger.as_ptr()).addresses[0].as_unchecked(), address) };
 
         const ADDRESS_COUNT: usize = 100;
         let mut prev = address;
-        for _ in 0..ADDRESS_COUNT {
+        for i in 0..ADDRESS_COUNT {
             let curr = rpc.get_new_address(None, None).unwrap();
 
             assert_ne!(prev, curr);
             assert!(curr.is_valid_for_network(Network::Regtest));
             assert!(!curr.is_valid_for_network(Network::Testnet));
             assert!(!curr.is_valid_for_network(Network::Signet));
             assert!(!curr.is_valid_for_network(Network::Bitcoin));
+            unsafe {
+                assert_eq!(
+                    *(*rpc.ledger.as_ptr()).addresses[i + 1].as_unchecked(),
+                    curr
+                )
+            };
 
             prev = curr;
         }

diff --git a/src/ledger/mod.rs b/src/ledger/mod.rs
@@ -5,11 +5,12 @@
 use bitcoin::{Address, TxOut};
 
 /// Mock Bitcoin ledger structure.
+#[derive(Clone)]
 pub struct Ledger {
     /// User's addresses.
-    addresses: Vec<Address>,
+    pub addresses: Vec<Address>,
     /// User's unspent transaction outputs.
-    utxos: Vec<TxOut>,
+    pub utxos: Vec<TxOut>,
 }
 
 impl Ledger {
@@ -26,8 +27,12 @@ impl Ledger {
     }
 
     /// Adds a new address for user.
-    pub fn add_address(&mut self, address: Address) {
-        self.addresses.push(address);
+    pub fn add_address(&self, address: Address) -> Self {
+        let mut ledger = self.clone().to_owned();
+
+        ledger.addresses.push(address);
+
+        ledger
     }
 }
 
@@ -56,6 +61,7 @@ mod tests {
         let mut ledger = Ledger::new();
 
         assert_eq!(ledger.utxos.len(), 0);
+        assert_eq!(ledger.addresses.len(), 0);
 
         // Generate a random address.
         let secp = Secp256k1::new();
@@ -66,7 +72,8 @@ mod tests {
         ])
         .unwrap();
         let address = Address::p2tr(&secp, xonly_public_key, None, Network::Regtest);
-        ledger.add_address(address);
+        ledger = ledger.add_address(address);
+        assert_eq!(ledger.addresses.len(), 1);
 
         // Insert a dummy UTXO.
         let utxo = TxOut {