From 2c0bc45ecf0241abc8bf9d72c8bbeded53d32a45 Mon Sep 17 00:00:00 2001
From: Wei Chen <wzc110@gmail.com>
Date: Thu, 8 Aug 2024 22:17:01 +0800
Subject: [PATCH] feat(testenv): Added `bdk_electrum` wait for Txid method

Added `wait_until_electrum_sees_txid` method to `TestEnv`. Both
`bdk_electrum` wait methods now have a `timeout` option. Removed
the exponential polling delay in lieu of a fixed delay inside the
`bdk_electrum` wait methods.
---
 crates/electrum/tests/test_electrum.rs |  6 ++--
 crates/testenv/src/lib.rs              | 41 +++++++++++++++++++++-----
 2 files changed, 38 insertions(+), 9 deletions(-)

diff --git a/crates/electrum/tests/test_electrum.rs b/crates/electrum/tests/test_electrum.rs
index 4abc1bbfc..afe50be0a 100644
--- a/crates/electrum/tests/test_electrum.rs
+++ b/crates/electrum/tests/test_electrum.rs
@@ -275,7 +275,9 @@ pub fn test_update_tx_graph_stop_gap() -> anyhow::Result<()> {
     Ok(())
 }
 
-/// Ensure that [`BdkElectrumClient::sync`] can confirm previously unconfirmed transactions in both reorg and no-reorg situations.
+/// 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 test_sync() -> anyhow::Result<()> {
     const SEND_AMOUNT: Amount = Amount::from_sat(10_000);
@@ -316,7 +318,7 @@ fn test_sync() -> anyhow::Result<()> {
         &mut recv_graph,
     )?;
 
-    // Check if balance is zero when transaction exists only in mempool.
+    // Check for unconfirmed balance when transaction exists only in mempool.
     assert_eq!(
         get_balance(&recv_chain, &recv_graph)?,
         Balance {
diff --git a/crates/testenv/src/lib.rs b/crates/testenv/src/lib.rs
index b0c75b30b..3c576cd14 100644
--- a/crates/testenv/src/lib.rs
+++ b/crates/testenv/src/lib.rs
@@ -168,22 +168,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`.
@@ -266,6 +292,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.
@@ -275,7 +302,7 @@ 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))
@@ -283,7 +310,7 @@ mod test {
 
         // 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))