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

2c0bc45 feat(testenv): Added `bdk_electrum` wait for Txid method (Wei Chen)
49e1a5e test(electrum): Test sync in reorg and no-reorg situations (Wei Chen)

Pull request description:

  <!-- You can erase any parts of this template not applicable to your Pull Request. -->
  Closes #1125.

  ### Description

  <!-- Describe the purpose of this PR, what's being adding and/or fixed -->
  Add new test for `bdk_electrum` to make sure previously unconfirmed transactions get confirmed again in both reorg and no-reorg situations.

  ### Changelog notice

  <!-- Notice the release manager should include in the release tag message changelog -->
  <!-- See https://keepachangelog.com/en/1.0.0/ for examples -->
  * Added `wait_until_electrum_sees_txid` method to `TestEnv`.
  * `wait_until_electrum_sees_block` now has a `Duration` input for timeout.
  * Removed exponential polling delay in `wait_until_electrum_sees_block`.
  * Added `test_sync` to `bdk_electrum` to make sure previously unconfirmed transactions get confirmed in both reorg and no-org situations.

  ### Checklists

  #### All Submissions:

  * [x] I've signed all my commits
  * [x] I followed the [contribution guidelines](https://github.com/bitcoindevkit/bdk/blob/master/CONTRIBUTING.md)
  * [x] I ran `cargo fmt` and `cargo clippy` before committing

ACKs for top commit:
  evanlinjin:
    ACK 2c0bc45

Tree-SHA512: 2218151267fe0a3dc3eefa1e7081902e25193f211b2624a3a9f2007b9b38eefb5b847754b9f8cbefba6be6a497cfbc85ac3e5a2db40f01fcc71a3c9d8427adf4

Author: evanlinjin

crates/electrum/tests/test_electrum.rs

@@ -1,15 +1,19 @@
 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 core::time::Duration;
 use std::collections::{BTreeSet, HashSet};
 use std::str::FromStr;
 
+// Batch size for `sync_with_electrum`.
+const BATCH_SIZE: usize = 5;
+
 fn get_balance(
     recv_chain: &LocalChain,
     recv_graph: &IndexedTxGraph<ConfirmationBlockTime, SpkTxOutIndex<()>>,
@@ -22,6 +26,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()?;
@@ -60,7 +97,7 @@ pub fn test_update_tx_graph_without_keychain() -> anyhow::Result<()> {
         None,
     )?;
     env.mine_blocks(1, None)?;
-    env.wait_until_electrum_sees_block()?;
+    env.wait_until_electrum_sees_block(Duration::from_secs(6))?;
 
     // use a full checkpoint linked list (since this is not what we are testing)
     let cp_tip = env.make_checkpoint_tip();
@@ -162,7 +199,7 @@ pub fn test_update_tx_graph_stop_gap() -> anyhow::Result<()> {
         None,
     )?;
     env.mine_blocks(1, None)?;
-    env.wait_until_electrum_sees_block()?;
+    env.wait_until_electrum_sees_block(Duration::from_secs(6))?;
 
     // use a full checkpoint linked list (since this is not what we are testing)
     let cp_tip = env.make_checkpoint_tip();
@@ -204,7 +241,7 @@ pub fn test_update_tx_graph_stop_gap() -> anyhow::Result<()> {
         None,
     )?;
     env.mine_blocks(1, None)?;
-    env.wait_until_electrum_sees_block()?;
+    env.wait_until_electrum_sees_block(Duration::from_secs(6))?;
 
     // A scan with gap limit 5 won't find the second transaction, but a scan with gap limit 6 will.
     // The last active indice won't be updated in the first case but will in the second one.
@@ -238,14 +275,11 @@ 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. After the transaction is confirmed after reorg, check if floating
+/// txouts for previous outputs were inserted for transaction fee calculation.
 #[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 +305,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(Duration::from_secs(6))?;
 
-    // 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, Duration::from_secs(6))?;
 
-    // Mine a block to confirm sent tx.
+    sync_with_electrum(
+        &client,
+        [spk_to_track.clone()],
+        &mut recv_chain,
+        &mut recv_graph,
+    )?;
+
+    // Check for unconfirmed balance 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(Duration::from_secs(6))?;
 
-    // Sync up to tip.
-    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 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(Duration::from_secs(6))?;
+
+    sync_with_electrum(
+        &client,
+        [spk_to_track.clone()],
+        &mut recv_chain,
+        &mut recv_graph,
+    )?;
 
-    // Check to see if tx is confirmed.
+    // 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(Duration::from_secs(6))?;
+
+    sync_with_electrum(&client, [spk_to_track], &mut recv_chain, &mut recv_graph)?;
+
+    // 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.
@@ -371,18 +458,14 @@ 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,
+    env.wait_until_electrum_sees_block(Duration::from_secs(6))?;
+    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();
@@ -407,24 +490,21 @@ fn tx_can_become_unconfirmed_after_reorg() -> anyhow::Result<()> {
     for depth in 1..=REORG_COUNT {
         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,
+        env.wait_until_electrum_sees_block(Duration::from_secs(6))?;
+        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

@@ -187,22 +187,48 @@ impl TestEnv {
     }
 
     /// This method waits for the Electrum notification indicating that a new block has been mined.
-    pub fn wait_until_electrum_sees_block(&self) -> anyhow::Result<()> {
+    /// `timeout` is the maximum [`Duration`] we want to wait for a response from Electrsd.
+    pub fn wait_until_electrum_sees_block(&self, timeout: Duration) -> anyhow::Result<()> {
         self.electrsd.client.block_headers_subscribe()?;
-        let mut delay = Duration::from_millis(64);
+        let delay = Duration::from_millis(200);
+        let start = std::time::Instant::now();
 
-        loop {
+        while start.elapsed() < timeout {
             self.electrsd.trigger()?;
             self.electrsd.client.ping()?;
             if self.electrsd.client.block_headers_pop()?.is_some() {
                 return Ok(());
             }
 
-            if delay.as_millis() < 512 {
-                delay = delay.mul_f32(2.0);
+            std::thread::sleep(delay);
+        }
+
+        Err(anyhow::Error::msg(
+            "Timed out waiting for Electrsd to get block header",
+        ))
+    }
+
+    /// This method waits for Electrsd to see a transaction with given `txid`. `timeout` is the
+    /// maximum [`Duration`] we want to wait for a response from Electrsd.
+    pub fn wait_until_electrum_sees_txid(
+        &self,
+        txid: Txid,
+        timeout: Duration,
+    ) -> anyhow::Result<()> {
+        let delay = Duration::from_millis(200);
+        let start = std::time::Instant::now();
+
+        while start.elapsed() < timeout {
+            if self.electrsd.client.transaction_get(&txid).is_ok() {
+                return Ok(());
             }
+
             std::thread::sleep(delay);
         }
+
+        Err(anyhow::Error::msg(
+            "Timed out waiting for Electrsd to get transaction",
+        ))
     }
 
     /// Invalidate a number of blocks of a given size `count`.
@@ -285,6 +311,7 @@ impl TestEnv {
 #[cfg(test)]
 mod test {
     use crate::TestEnv;
+    use core::time::Duration;
     use electrsd::bitcoind::{anyhow::Result, bitcoincore_rpc::RpcApi};
 
     /// This checks that reorgs initiated by `bitcoind` is detected by our `electrsd` instance.
@@ -294,15 +321,15 @@ mod test {
 
         // Mine some blocks.
         env.mine_blocks(101, None)?;
-        env.wait_until_electrum_sees_block()?;
+        env.wait_until_electrum_sees_block(Duration::from_secs(6))?;
         let height = env.bitcoind.client.get_block_count()?;
         let blocks = (0..=height)
             .map(|i| env.bitcoind.client.get_block_hash(i))
             .collect::<Result<Vec<_>, _>>()?;
 
         // Perform reorg on six blocks.
         env.reorg(6)?;
-        env.wait_until_electrum_sees_block()?;
+        env.wait_until_electrum_sees_block(Duration::from_secs(6))?;
         let reorged_height = env.bitcoind.client.get_block_count()?;
         let reorged_blocks = (0..=height)
             .map(|i| env.bitcoind.client.get_block_hash(i))