diff --git a/swap/src/bitcoin.rs b/swap/src/bitcoin.rs index c44facdd..09ba21cd 100644 --- a/swap/src/bitcoin.rs +++ b/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 bob_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 = 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); diff --git a/swap/src/bitcoin/lock.rs b/swap/src/bitcoin/lock.rs index 2f82ff77..559ef1e4 100644 --- a/swap/src/bitcoin/lock.rs +++ b/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(); diff --git a/swap/src/bitcoin/wallet.rs b/swap/src/bitcoin/wallet.rs index a45e1a84..b544e650 100644 --- a/swap/src/bitcoin/wallet.rs +++ b/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 { - Self::new_funded(amount, 1.0, 1000) + pub fn new(amount: u64) -> Self { + 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. - /// - /// This will create invalid transactions but can be useful if you want full - /// control over the output amounts. - pub fn new_funded_zero_fees(amount: u64) -> Self { - Self::new_funded(amount, 0.0, 0) + pub fn with_zero_fees(self) -> Self { + Self { + sats_per_vb: 0.0, + min_relay_fee_sats: 0, + ..self + } + } + + pub fn with_fees(self, sats_per_vb: f32, min_relay_fee_sats: u64) -> Self { + Self { + sats_per_vb, + min_relay_fee_sats, + ..self + } } - /// Creates a new, funded wallet to be used within tests. - pub fn new_funded(amount: u64, sats_per_vb: f32, min_relay_fee_sats: u64) -> 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, - testutils! { - @tx ( (@external descriptors, 0) => amount ) (@confirmations 1) - }, - Some(100) - ); + + for index in 0..self.num_utxos { + bdk::populate_test_db!( + &mut database, + testutils! { + @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), - min_relay_fee: bitcoin::Amount::from_sat(min_relay_fee_sats), + fee_rate: FeeRate::from_sat_per_vb(self.sats_per_vb), + 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::
() .unwrap();