Introduce WalletBuilder for creating test instances of wallet

swap/src/bitcoin.rs

@@ -21,6 +21,9 @@ pub use ecdsa_fun::fun::Scalar;
 pub use ecdsa_fun::Signature;
 pub use wallet::Wallet;
 
+#[cfg(test)]
+pub use wallet::WalletBuilder;
+
 use crate::bitcoin::wallet::ScriptStatus;
 use ::bitcoin::hashes::hex::ToHex;
 use ::bitcoin::hashes::Hash;
@@ -317,8 +320,8 @@ mod tests {
 
     #[tokio::test]
     async fn calculate_transaction_weights() {
-        let alice_wallet = Wallet::new_funded_default_fees(Amount::ONE_BTC.as_sat());
+        let alice_wallet = WalletBuilder::new(Amount::ONE_BTC.as_sat()).build();
-        let bob_wallet = Wallet::new_funded_default_fees(Amount::ONE_BTC.as_sat());
+        let bob_wallet = WalletBuilder::new(Amount::ONE_BTC.as_sat()).build();
         let spending_fee = Amount::from_sat(1_000);
         let btc_amount = Amount::from_sat(500_000);
         let xmr_amount = crate::monero::Amount::from_piconero(10000);

swap/src/bitcoin/lock.rs

@@ -183,11 +183,12 @@ impl Watchable for TxLock {
 mod tests {
     use super::*;
     use crate::bitcoin::wallet::StaticFeeRate;
+    use crate::bitcoin::WalletBuilder;
 
     #[tokio::test]
     async fn given_bob_sends_good_psbt_when_reconstructing_then_succeeeds() {
         let (A, B) = alice_and_bob();
-        let wallet = Wallet::new_funded_default_fees(50000);
+        let wallet = WalletBuilder::new(50_000).build();
         let agreed_amount = Amount::from_sat(10000);
 
         let psbt = bob_make_psbt(A, B, &wallet, agreed_amount).await;
@@ -201,7 +202,7 @@ mod tests {
         let (A, B) = alice_and_bob();
         let fees = 610;
         let agreed_amount = Amount::from_sat(10000);
-        let wallet = Wallet::new_funded_default_fees(agreed_amount.as_sat() + fees);
+        let wallet = WalletBuilder::new(agreed_amount.as_sat() + fees).build();
 
         let psbt = bob_make_psbt(A, B, &wallet, agreed_amount).await;
         assert_eq!(
@@ -217,7 +218,7 @@ mod tests {
     #[tokio::test]
     async fn given_bob_is_sending_less_than_agreed_when_reconstructing_txlock_then_fails() {
         let (A, B) = alice_and_bob();
-        let wallet = Wallet::new_funded_default_fees(50000);
+        let wallet = WalletBuilder::new(50_000).build();
         let agreed_amount = Amount::from_sat(10000);
 
         let bad_amount = Amount::from_sat(5000);
@@ -230,7 +231,7 @@ mod tests {
     #[tokio::test]
     async fn given_bob_is_sending_to_a_bad_output_reconstructing_txlock_then_fails() {
         let (A, B) = alice_and_bob();
-        let wallet = Wallet::new_funded_default_fees(50000);
+        let wallet = WalletBuilder::new(50_000).build();
         let agreed_amount = Amount::from_sat(10000);
 
         let E = eve();

swap/src/bitcoin/wallet.rs

@@ -547,48 +547,84 @@ impl EstimateFeeRate for StaticFeeRate {
 }
 
 #[cfg(test)]
-impl Wallet<(), bdk::database::MemoryDatabase, StaticFeeRate> {
+pub struct WalletBuilder {
+    utxo_amount: u64,
+    sats_per_vb: f32,
+    min_relay_fee_sats: u64,
+    key: bitcoin::util::bip32::ExtendedPrivKey,
+    num_utxos: u8,
+}
+
+#[cfg(test)]
+impl WalletBuilder {
     /// Creates a new, funded wallet with sane default fees.
     ///
     /// Unless you are testing things related to fees, this is likely what you
     /// want.
-    pub fn new_funded_default_fees(amount: u64) -> Self {
+    pub fn new(amount: u64) -> Self {
-        Self::new_funded(amount, 1.0, 1000)
+        WalletBuilder {
+            utxo_amount: amount,
+            sats_per_vb: 1.0,
+            min_relay_fee_sats: 1000,
+            key: "tprv8ZgxMBicQKsPeZRHk4rTG6orPS2CRNFX3njhUXx5vj9qGog5ZMH4uGReDWN5kCkY3jmWEtWause41CDvBRXD1shKknAMKxT99o9qUTRVC6m".parse().unwrap(),
+            num_utxos: 1,
+        }
     }
 
-    /// Creates a new, funded wallet that doesn't pay any fees.
+    pub fn with_zero_fees(self) -> Self {
-    ///
+        Self {
-    /// This will create invalid transactions but can be useful if you want full
+            sats_per_vb: 0.0,
-    /// control over the output amounts.
+            min_relay_fee_sats: 0,
-    pub fn new_funded_zero_fees(amount: u64) -> Self {
+            ..self
-        Self::new_funded(amount, 0.0, 0)
+        }
     }
 
-    /// Creates a new, funded wallet to be used within tests.
+    pub fn with_fees(self, sats_per_vb: f32, min_relay_fee_sats: u64) -> Self {
-    pub fn new_funded(amount: u64, sats_per_vb: f32, min_relay_fee_sats: u64) -> Self {
+        Self {
+            sats_per_vb,
+            min_relay_fee_sats,
+            ..self
+        }
+    }
+
+    pub fn with_key(self, key: bitcoin::util::bip32::ExtendedPrivKey) -> Self {
+        Self { key, ..self }
+    }
+
+    pub fn with_num_utxos(self, number: u8) -> Self {
+        Self {
+            num_utxos: number,
+            ..self
+        }
+    }
+
+    pub fn build(self) -> Wallet<(), bdk::database::MemoryDatabase, StaticFeeRate> {
         use bdk::database::MemoryDatabase;
         use bdk::{LocalUtxo, TransactionDetails};
         use bitcoin::OutPoint;
         use testutils::testutils;
 
-        let descriptors = testutils!(@descriptors ("wpkh(tprv8ZgxMBicQKsPeZRHk4rTG6orPS2CRNFX3njhUXx5vj9qGog5ZMH4uGReDWN5kCkY3jmWEtWause41CDvBRXD1shKknAMKxT99o9qUTRVC6m/*)"));
+        let descriptors = testutils!(@descriptors (&format!("wpkh({}/*)", self.key)));
 
         let mut database = MemoryDatabase::new();
-        bdk::populate_test_db!(
+
-            &mut database,
+        for index in 0..self.num_utxos {
-            testutils! {
+            bdk::populate_test_db!(
-                @tx ( (@external descriptors, 0) => amount ) (@confirmations 1)
+                &mut database,
-            },
+                testutils! {
-            Some(100)
+                    @tx ( (@external descriptors, index as u32) => self.utxo_amount ) (@confirmations 1)
-        );
+                },
+                Some(100)
+            );
+        }
 
         let wallet =
             bdk::Wallet::new_offline(&descriptors.0, None, Network::Regtest, database).unwrap();
 
-        Self {
+        Wallet {
             client: Arc::new(Mutex::new(StaticFeeRate {
-                fee_rate: FeeRate::from_sat_per_vb(sats_per_vb),
+                fee_rate: FeeRate::from_sat_per_vb(self.sats_per_vb),
-                min_relay_fee: bitcoin::Amount::from_sat(min_relay_fee_sats),
+                min_relay_fee: bitcoin::Amount::from_sat(self.min_relay_fee_sats),
             })),
             wallet: Arc::new(Mutex::new(wallet)),
             finality_confirmations: 1,
@@ -1047,7 +1083,7 @@ mod tests {
 
     #[tokio::test]
     async fn given_no_balance_returns_amount_0() {
-        let wallet = Wallet::new_funded(0, 1.0, 1);
+        let wallet = WalletBuilder::new(0).with_fees(1.0, 1).build();
         let amount = wallet.max_giveable(TxLock::script_size()).await.unwrap();
 
         assert_eq!(amount, Amount::ZERO);
@@ -1055,7 +1091,7 @@ mod tests {
 
     #[tokio::test]
     async fn given_balance_below_min_relay_fee_returns_amount_0() {
-        let wallet = Wallet::new_funded(1000, 1.0, 1001);
+        let wallet = WalletBuilder::new(1000).with_fees(1.0, 1001).build();
         let amount = wallet.max_giveable(TxLock::script_size()).await.unwrap();
 
         assert_eq!(amount, Amount::ZERO);
@@ -1063,7 +1099,7 @@ mod tests {
 
     #[tokio::test]
     async fn given_balance_above_relay_fee_returns_amount_greater_0() {
-        let wallet = Wallet::new_funded_default_fees(10_000);
+        let wallet = WalletBuilder::new(10_000).build();
         let amount = wallet.max_giveable(TxLock::script_size()).await.unwrap();
 
         assert!(amount.as_sat() > 0);
@@ -1083,7 +1119,7 @@ mod tests {
         let balance = 2000;
 
         // We don't care about fees in this test, thus use a zero fee rate
-        let wallet = Wallet::new_funded_zero_fees(balance);
+        let wallet = WalletBuilder::new(balance).with_zero_fees().build();
 
         // sorting is only relevant for amounts that have a change output
         // if the change output is below dust it will be dropped by the BDK
@@ -1108,7 +1144,7 @@ mod tests {
 
     #[tokio::test]
     async fn can_override_change_address() {
-        let wallet = Wallet::new_funded_default_fees(50_000);
+        let wallet = WalletBuilder::new(50_000).build();
         let custom_change = "bcrt1q08pfqpsyrt7acllzyjm8q5qsz5capvyahm49rw"
             .parse::<Address>()
             .unwrap();