test(electrum): Test sync in reorg and no-reorg situations

crates/electrum/tests/test_electrum.rs

@@ -1,15 +1,18 @@
 use bdk_chain::{
     bitcoin::{hashes::Hash, Address, Amount, ScriptBuf, Txid, WScriptHash},
     local_chain::LocalChain,
-    spk_client::{FullScanRequest, SyncRequest},
+    spk_client::{FullScanRequest, SyncRequest, SyncResult},
     spk_txout::SpkTxOutIndex,
-    Balance, ConfirmationBlockTime, IndexedTxGraph,
+    Balance, ConfirmationBlockTime, IndexedTxGraph, Indexer, Merge,
 };
 use bdk_electrum::BdkElectrumClient;
 use bdk_testenv::{anyhow, bitcoincore_rpc::RpcApi, TestEnv};
 use std::collections::{BTreeSet, HashSet};
 use std::str::FromStr;
 
+// Batch size for `electrum_sync`.
+const BATCH_SIZE: usize = 5;
+
 fn get_balance(
     recv_chain: &LocalChain,
     recv_graph: &IndexedTxGraph<ConfirmationBlockTime, SpkTxOutIndex<()>>,
@@ -22,6 +25,39 @@ fn get_balance(
     Ok(balance)
 }
 
+fn sync_with_electrum<I, Spks>(
+    client: &BdkElectrumClient<electrum_client::Client>,
+    spks: Spks,
+    chain: &mut LocalChain,
+    graph: &mut IndexedTxGraph<ConfirmationBlockTime, I>,
+) -> anyhow::Result<SyncResult>
+where
+    I: Indexer,
+    I::ChangeSet: Default + Merge,
+    Spks: IntoIterator<Item = ScriptBuf>,
+    Spks::IntoIter: ExactSizeIterator + Send + 'static,
+{
+    let mut update = client.sync(
+        SyncRequest::from_chain_tip(chain.tip()).chain_spks(spks),
+        BATCH_SIZE,
+        true,
+    )?;
+
+    // Update `last_seen` to be able to calculate balance for unconfirmed transactions.
+    let now = std::time::UNIX_EPOCH
+        .elapsed()
+        .expect("must get time")
+        .as_secs();
+    let _ = update.graph_update.update_last_seen_unconfirmed(now);
+
+    let _ = chain
+        .apply_update(update.chain_update.clone())
+        .map_err(|err| anyhow::anyhow!("LocalChain update error: {:?}", err))?;
+    let _ = graph.apply_update(update.graph_update.clone());
+
+    Ok(update)
+}
+
 #[test]
 pub fn test_update_tx_graph_without_keychain() -> anyhow::Result<()> {
     let env = TestEnv::new()?;
@@ -238,14 +274,9 @@ pub fn test_update_tx_graph_stop_gap() -> anyhow::Result<()> {
     Ok(())
 }
 
-/// Ensure that [`ElectrumExt`] can sync properly.
-///
-/// 1. Mine 101 blocks.
-/// 2. Send a tx.
-/// 3. Mine extra block to confirm sent tx.
-/// 4. Check [`Balance`] to ensure tx is confirmed.
+/// Ensure that [`BdkElectrumClient::sync`] can confirm previously unconfirmed transactions in both reorg and no-reorg situations.
 #[test]
-fn scan_detects_confirmed_tx() -> anyhow::Result<()> {
+fn test_sync() -> anyhow::Result<()> {
     const SEND_AMOUNT: Amount = Amount::from_sat(10_000);
 
     let env = TestEnv::new()?;
@@ -271,35 +302,88 @@ fn scan_detects_confirmed_tx() -> anyhow::Result<()> {
 
     // Mine some blocks.
     env.mine_blocks(101, Some(addr_to_mine))?;
+    env.wait_until_electrum_sees_block()?;
 
-    // Create transaction that is tracked by our receiver.
-    env.send(&addr_to_track, SEND_AMOUNT)?;
+    // Broadcast transaction to mempool.
+    let txid = env.send(&addr_to_track, SEND_AMOUNT)?;
+    env.wait_until_electrum_sees_txid(txid)?;
 
-    // Mine a block to confirm sent tx.
+    sync_with_electrum(
+        &client,
+        [spk_to_track.clone()],
+        &mut recv_chain,
+        &mut recv_graph,
+    )?;
+
+    // Check if balance is zero when transaction exists only in mempool.
+    assert_eq!(
+        get_balance(&recv_chain, &recv_graph)?,
+        Balance {
+            trusted_pending: SEND_AMOUNT,
+            ..Balance::default()
+        },
+        "balance must be correct",
+    );
+
+    // Mine block to confirm transaction.
     env.mine_blocks(1, None)?;
+    env.wait_until_electrum_sees_block()?;
 
-    // Sync up to tip.
+    sync_with_electrum(
+        &client,
+        [spk_to_track.clone()],
+        &mut recv_chain,
+        &mut recv_graph,
+    )?;
+
+    // Check if balance is correct when transaction is confirmed.
+    assert_eq!(
+        get_balance(&recv_chain, &recv_graph)?,
+        Balance {
+            confirmed: SEND_AMOUNT,
+            ..Balance::default()
+        },
+        "balance must be correct",
+    );
+
+    // Perform reorg on block with confirmed transaction.
+    env.reorg_empty_blocks(1)?;
     env.wait_until_electrum_sees_block()?;
-    let update = client.sync(
-        SyncRequest::from_chain_tip(recv_chain.tip()).chain_spks(core::iter::once(spk_to_track)),
-        5,
-        true,
+
+    sync_with_electrum(
+        &client,
+        [spk_to_track.clone()],
+        &mut recv_chain,
+        &mut recv_graph,
     )?;
 
-    let _ = recv_chain
-        .apply_update(update.chain_update)
-        .map_err(|err| anyhow::anyhow!("LocalChain update error: {:?}", err))?;
-    let _ = recv_graph.apply_update(update.graph_update);
+    // Check if balance is correct when transaction returns to mempool.
+    assert_eq!(
+        get_balance(&recv_chain, &recv_graph)?,
+        Balance {
+            trusted_pending: SEND_AMOUNT,
+            ..Balance::default()
+        },
+    );
+
+    // Mine block to confirm transaction again.
+    env.mine_blocks(1, None)?;
+    env.wait_until_electrum_sees_block()?;
+
+    sync_with_electrum(&client, [spk_to_track], &mut recv_chain, &mut recv_graph)?;
 
-    // Check to see if tx is confirmed.
+    // Check if balance is correct once transaction is confirmed again.
     assert_eq!(
         get_balance(&recv_chain, &recv_graph)?,
         Balance {
             confirmed: SEND_AMOUNT,
             ..Balance::default()
         },
+        "balance must be correct",
     );
 
+    // Check to see if we have the floating txouts available from our transactions' previous outputs
+    // in order to calculate transaction fees.
     for tx in recv_graph.graph().full_txs() {
         // Retrieve the calculated fee from `TxGraph`, which will panic if we do not have the
         // floating txouts available from the transaction's previous outputs.
@@ -372,17 +456,13 @@ fn tx_can_become_unconfirmed_after_reorg() -> anyhow::Result<()> {
 
     // Sync up to tip.
     env.wait_until_electrum_sees_block()?;
-    let update = client.sync(
-        SyncRequest::from_chain_tip(recv_chain.tip()).chain_spks([spk_to_track.clone()]),
-        5,
-        false,
+    let update = sync_with_electrum(
+        &client,
+        [spk_to_track.clone()],
+        &mut recv_chain,
+        &mut recv_graph,
     )?;
 
-    let _ = recv_chain
-        .apply_update(update.chain_update)
-        .map_err(|err| anyhow::anyhow!("LocalChain update error: {:?}", err))?;
-    let _ = recv_graph.apply_update(update.graph_update.clone());
-
     // Retain a snapshot of all anchors before reorg process.
     let initial_anchors = update.graph_update.all_anchors();
     let anchors: Vec<_> = initial_anchors.iter().cloned().collect();
@@ -408,23 +488,20 @@ fn tx_can_become_unconfirmed_after_reorg() -> anyhow::Result<()> {
         env.reorg_empty_blocks(depth)?;
 
         env.wait_until_electrum_sees_block()?;
-        let update = client.sync(
-            SyncRequest::from_chain_tip(recv_chain.tip()).chain_spks([spk_to_track.clone()]),
-            5,
-            false,
+        let update = sync_with_electrum(
+            &client,
+            [spk_to_track.clone()],
+            &mut recv_chain,
+            &mut recv_graph,
         )?;
 
-        let _ = recv_chain
-            .apply_update(update.chain_update)
-            .map_err(|err| anyhow::anyhow!("LocalChain update error: {:?}", err))?;
-
         // Check that no new anchors are added during current reorg.
         assert!(initial_anchors.is_superset(update.graph_update.all_anchors()));
-        let _ = recv_graph.apply_update(update.graph_update);
 
         assert_eq!(
             get_balance(&recv_chain, &recv_graph)?,
             Balance {
+                trusted_pending: SEND_AMOUNT * depth as u64,
                 confirmed: SEND_AMOUNT * (REORG_COUNT - depth) as u64,
                 ..Balance::default()
             },

crates/testenv/src/lib.rs

@@ -186,6 +186,22 @@ impl TestEnv {
         }
     }
 
+    /// This method waits for Electrum to see a transaction with given `txid`.
+    pub fn wait_until_electrum_sees_txid(&self, txid: Txid) -> anyhow::Result<()> {
+        let mut delay = Duration::from_millis(64);
+
+        loop {
+            if self.electrsd.client.transaction_get(&txid).is_ok() {
+                return Ok(());
+            }
+
+            if delay.as_millis() < 512 {
+                delay = delay.mul_f32(2.0);
+            }
+            std::thread::sleep(delay);
+        }
+    }
+
     /// Invalidate a number of blocks of a given size `count`.
     pub fn invalidate_blocks(&self, count: usize) -> anyhow::Result<()> {
         let mut hash = self.bitcoind.client.get_best_block_hash()?;