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

crates/electrum/tests/test_electrum.rs

@@ -1,12 +1,13 @@
 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 electrum_client::Client;
 use std::collections::{BTreeSet, HashSet};
 use std::str::FromStr;
 
@@ -22,10 +23,34 @@ fn get_balance(
     Ok(balance)
 }
 
+fn electrum_sync<I>(
+    client: &BdkElectrumClient<Client>,
+    spks: &Vec<ScriptBuf>,
+    batch_size: usize,
+    chain: &mut LocalChain,
+    graph: &mut IndexedTxGraph<ConfirmationBlockTime, I>,
+) -> anyhow::Result<SyncResult>
+where
+    I: Indexer,
+    I::ChangeSet: Default + Merge,
+{
+    let update = client.sync(
+        SyncRequest::from_chain_tip(chain.tip()).chain_spks(spks.clone()),
+        batch_size,
+        true,
+    )?;
+    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()?;
-    let electrum_client = electrum_client::Client::new(env.electrsd.electrum_url.as_str())?;
+    let electrum_client = Client::new(env.electrsd.electrum_url.as_str())?;
     let client = BdkElectrumClient::new(electrum_client);
 
     let receive_address0 =
@@ -123,7 +148,7 @@ pub fn test_update_tx_graph_without_keychain() -> anyhow::Result<()> {
 #[test]
 pub fn test_update_tx_graph_stop_gap() -> anyhow::Result<()> {
     let env = TestEnv::new()?;
-    let electrum_client = electrum_client::Client::new(env.electrsd.electrum_url.as_str())?;
+    let electrum_client = Client::new(env.electrsd.electrum_url.as_str())?;
     let client = BdkElectrumClient::new(electrum_client);
     let _block_hashes = env.mine_blocks(101, None)?;
 
@@ -238,18 +263,13 @@ 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`] can sync properly 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()?;
-    let electrum_client = electrum_client::Client::new(env.electrsd.electrum_url.as_str())?;
+    let electrum_client = Client::new(env.electrsd.electrum_url.as_str())?;
     let client = BdkElectrumClient::new(electrum_client);
 
     // Setup addresses.
@@ -272,32 +292,91 @@ fn scan_detects_confirmed_tx() -> anyhow::Result<()> {
     // Mine some blocks.
     env.mine_blocks(101, Some(addr_to_mine))?;
 
-    // Create transaction that is tracked by our receiver.
+    // Broadcast transaction to mempool.
     env.send(&addr_to_track, SEND_AMOUNT)?;
+    env.wait_until_electrum_sees_block()?;
+
+    electrum_sync(
+        &client,
+        &vec![spk_to_track.clone()],
+        5,
+        &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 {
+            confirmed: Amount::from_sat(0),
+            ..Balance::default()
+        },
+        "initial balance must be correct",
+    );
 
-    // Mine a block to confirm sent tx.
+    // Mine block to confirm transaction.
     env.mine_blocks(1, None)?;
+    env.wait_until_electrum_sees_block()?;
 
-    // Sync up to tip.
+    electrum_sync(
+        &client,
+        &vec![spk_to_track.clone()],
+        5,
+        &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()
+        },
+        "initial 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)),
+
+    electrum_sync(
+        &client,
+        &vec![spk_to_track.clone()],
         5,
-        true,
+        &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 {
+            confirmed: Amount::from_sat(0),
+            ..Balance::default()
+        },
+    );
+
+    // Mine block to confirm transaction again.
+    env.mine_blocks(1, None)?;
+    env.wait_until_electrum_sees_block()?;
+
+    electrum_sync(
+        &client,
+        &vec![spk_to_track],
+        5,
+        &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()
         },
+        "initial balance must be correct",
     );
 
     for tx in recv_graph.graph().full_txs() {
@@ -339,7 +418,7 @@ fn tx_can_become_unconfirmed_after_reorg() -> anyhow::Result<()> {
     const SEND_AMOUNT: Amount = Amount::from_sat(10_000);
 
     let env = TestEnv::new()?;
-    let electrum_client = electrum_client::Client::new(env.electrsd.electrum_url.as_str())?;
+    let electrum_client = Client::new(env.electrsd.electrum_url.as_str())?;
     let client = BdkElectrumClient::new(electrum_client);
 
     // Setup addresses.
@@ -372,17 +451,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()]),
+    let update = electrum_sync(
+        &client,
+        &vec![spk_to_track.clone()],
         5,
-        false,
+        &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,19 +484,16 @@ 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()]),
+        let update = electrum_sync(
+            &client,
+            &vec![spk_to_track.clone()],
             5,
-            false,
+            &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)?,