Refactor punish test and punish assertions

swap/tests/punish.rs

@@ -1,21 +1,5 @@
-use crate::testutils::{init_alice, init_bob};
-use futures::{
-    future::{join, select, Either},
-    FutureExt,
-};
-use get_port::get_port;
-use libp2p::Multiaddr;
 use rand::rngs::OsRng;
-use swap::{
+use swap::protocol::{alice, bob, bob::BobState};
-    bitcoin,
-    config::Config,
-    monero,
-    protocol::{alice, alice::AliceState, bob, bob::BobState},
-    seed::Seed,
-};
-use testcontainers::clients::Cli;
-use testutils::init_tracing;
-use uuid::Uuid;
 
 pub mod testutils;
 
@@ -23,128 +7,64 @@ pub mod testutils;
 /// the encsig and fail to refund or redeem. Alice punishes.
 #[tokio::test]
 async fn alice_punishes_if_bob_never_acts_after_fund() {
-    let _guard = init_tracing();
+    testutils::test(|alice_harness, bob_harness| async move {
+        let alice = alice_harness.new_alice().await;
+        let bob = bob_harness.new_bob().await;
 
-    let cli = Cli::default();
+        let alice_swap = alice::swap(
-    let (
+            alice.state,
-        monero,
+            alice.event_loop_handle,
-        testutils::Containers {
+            alice.bitcoin_wallet.clone(),
-            bitcoind,
+            alice.monero_wallet.clone(),
-            monerods: _monerods,
+            alice.config,
-        },
+            alice.swap_id,
-    ) = testutils::init_containers(&cli).await;
+            alice.db,
+        );
+        let alice_swap_handle = tokio::spawn(alice_swap);
 
-    let btc_to_swap = bitcoin::Amount::from_sat(1_000_000);
+        let bob_state = bob::run_until(
-    let xmr_to_swap = monero::Amount::from_piconero(1_000_000_000_000);
+            bob.state,
-
+            bob::swap::is_btc_locked,
-    let bob_btc_starting_balance = btc_to_swap * 10;
+            bob.event_loop_handle,
-
+            bob.db,
-    let alice_btc_starting_balance = bitcoin::Amount::ZERO;
+            bob.bitcoin_wallet.clone(),
-    let alice_xmr_starting_balance = xmr_to_swap * 10;
+            bob.monero_wallet.clone(),
-
+            OsRng,
-    let port = get_port().expect("Failed to find a free port");
+            bob.swap_id,
-    let alice_multiaddr: Multiaddr = format!("/ip4/127.0.0.1/tcp/{}", port)
-        .parse()
-        .expect("failed to parse Alice's address");
-
-    let config = Config::regtest();
-
-    let (
-        alice_state,
-        mut alice_event_loop,
-        alice_event_loop_handle,
-        alice_btc_wallet,
-        alice_xmr_wallet,
-        alice_db,
-    ) = init_alice(
-        &bitcoind,
-        &monero,
-        btc_to_swap,
-        xmr_to_swap,
-        alice_xmr_starting_balance,
-        alice_multiaddr.clone(),
-        config,
-        Seed::random().unwrap(),
-    )
-    .await;
-
-    let (bob_state, bob_event_loop, bob_event_loop_handle, bob_btc_wallet, bob_xmr_wallet, bob_db) =
-        init_bob(
-            alice_multiaddr,
-            alice_event_loop.peer_id(),
-            &bitcoind,
-            &monero,
-            btc_to_swap,
-            bob_btc_starting_balance,
-            xmr_to_swap,
-            config,
         )
-        .await;
-
-    let bob_btc_locked_fut = bob::swap::run_until(
-        bob_state,
-        bob::swap::is_btc_locked,
-        bob_event_loop_handle,
-        bob_db,
-        bob_btc_wallet.clone(),
-        bob_xmr_wallet.clone(),
-        OsRng,
-        Uuid::new_v4(),
-    )
-    .boxed();
-
-    let bob_fut = select(bob_btc_locked_fut, bob_event_loop.run().boxed());
-
-    let alice_fut = alice::swap::swap(
-        alice_state,
-        alice_event_loop_handle,
-        alice_btc_wallet.clone(),
-        alice_xmr_wallet.clone(),
-        Config::regtest(),
-        Uuid::new_v4(),
-        alice_db,
-    )
-    .boxed();
-
-    let alice_fut = select(alice_fut, alice_event_loop.run().boxed());
-
-    // Wait until alice has locked xmr and bob has locked btc
-    let (alice_state, bob_state) = join(alice_fut, bob_fut).await;
-
-    let alice_state = match alice_state {
-        Either::Left((state, _)) => state.unwrap(),
-        Either::Right(_) => panic!("Alice event loop should not terminate."),
-    };
-
-    let bob_state = match bob_state {
-        Either::Left((state, _)) => state.unwrap(),
-        Either::Right(_) => panic!("Bob event loop should not terminate."),
-    };
-
-    assert!(matches!(alice_state, AliceState::BtcPunished));
-    let bob_state3 = if let BobState::BtcLocked(state3, ..) = bob_state {
-        state3
-    } else {
-        panic!("Bob in unexpected state");
-    };
-
-    let btc_alice_final = alice_btc_wallet.as_ref().balance().await.unwrap();
-    let btc_bob_final = bob_btc_wallet.as_ref().balance().await.unwrap();
-
-    // lock_tx_bitcoin_fee is determined by the wallet, it is not necessarily equal
-    // to TX_FEE
-    let lock_tx_bitcoin_fee = bob_btc_wallet
-        .transaction_fee(bob_state3.tx_lock_id())
         .await
         .unwrap();
 
-    assert_eq!(
+        assert!(matches!(bob_state, BobState::BtcLocked {..}));
-        btc_alice_final,
-        alice_btc_starting_balance + btc_to_swap - bitcoin::Amount::from_sat(2 * bitcoin::TX_FEE)
-    );
 
-    assert_eq!(
+        let alice_state = alice_swap_handle.await.unwrap();
-        btc_bob_final,
+        alice_harness.assert_punished(alice_state.unwrap()).await;
-        bob_btc_starting_balance - btc_to_swap - lock_tx_bitcoin_fee
+
-    );
+        // Restart Bob after Alice punished to ensure Bob transitions to
+        // punished and does not run indefinitely
+        let bob = bob_harness.recover_bob_from_db().await;
+        assert!(matches!(bob.state, BobState::BtcLocked {..}));
+
+        // TODO: make lock-tx-id available in final states
+        let lock_tx_id = if let BobState::BtcLocked(state3) = bob_state {
+            state3.tx_lock_id()
+        } else {
+            panic!("Bob in unexpected state");
+        };
+
+        let bob_state = bob::swap(
+            bob.state,
+            bob.event_loop_handle,
+            bob.db,
+            bob.bitcoin_wallet.clone(),
+            bob.monero_wallet.clone(),
+            OsRng,
+            bob.swap_id,
+        )
+        .await
+        .unwrap();
+
+        bob_harness.assert_punished(bob_state, lock_tx_id).await;
+    })
+    .await;
 }

swap/tests/testutils/mod.rs

@@ -65,16 +65,29 @@ impl AliceHarness {
     pub async fn assert_redeemed(&self, state: AliceState) {
         assert!(matches!(state, AliceState::BtcRedeemed));
 
-        let btc_alice_final = self.bitcoin_wallet.as_ref().balance().await.unwrap();
+        let btc_balance_after_swap = self.bitcoin_wallet.as_ref().balance().await.unwrap();
-
         assert_eq!(
-            btc_alice_final,
+            btc_balance_after_swap,
             self.starting_balances.btc + self.swap_amounts.btc
                 - bitcoin::Amount::from_sat(bitcoin::TX_FEE)
         );
 
-        let xmr_alice_final = self.monero_wallet.as_ref().get_balance().await.unwrap();
+        let xmr_balance_after_swap = self.monero_wallet.as_ref().get_balance().await.unwrap();
-        assert!(xmr_alice_final <= self.starting_balances.xmr - self.swap_amounts.xmr);
+        assert!(xmr_balance_after_swap <= self.starting_balances.xmr - self.swap_amounts.xmr);
+    }
+
+    pub async fn assert_punished(&self, state: AliceState) {
+        assert!(matches!(state, AliceState::BtcPunished));
+
+        let btc_balance_after_swap = self.bitcoin_wallet.as_ref().balance().await.unwrap();
+        assert_eq!(
+            btc_balance_after_swap,
+            self.starting_balances.btc + self.swap_amounts.btc
+                - bitcoin::Amount::from_sat(2 * bitcoin::TX_FEE)
+        );
+
+        let xnr_balance_after_swap = self.monero_wallet.as_ref().get_balance().await.unwrap();
+        assert!(xnr_balance_after_swap <= self.starting_balances.xmr - self.swap_amounts.xmr);
     }
 
     pub async fn new(
@@ -294,15 +307,38 @@ impl BobHarness {
     pub async fn assert_redeemed(&self, state: BobState) {
         assert!(matches!(state, BobState::XmrRedeemed));
 
-        let btc_bob_final = self.bitcoin_wallet.as_ref().balance().await.unwrap();
+        let btc_balance_after_swap = self.bitcoin_wallet.as_ref().balance().await.unwrap();
+        assert!(btc_balance_after_swap <= self.starting_balances.btc - self.swap_amounts.btc);
+
+        // Ensure that Bob's balance is refreshed as we use a newly created wallet
         self.monero_wallet.as_ref().inner.refresh().await.unwrap();
-        let xmr_bob_final = self.monero_wallet.as_ref().get_balance().await.unwrap();
+        let xmr_balance_after_swap = self.monero_wallet.as_ref().get_balance().await.unwrap();
-        assert!(btc_bob_final <= self.starting_balances.btc - self.swap_amounts.btc);
         assert_eq!(
-            xmr_bob_final,
+            xmr_balance_after_swap,
             self.starting_balances.xmr + self.swap_amounts.xmr
         );
     }
+
+    pub async fn assert_punished(&self, state: BobState, lock_tx_id: ::bitcoin::Txid) {
+        assert!(matches!(state, BobState::BtcPunished));
+
+        // lock_tx_bitcoin_fee is determined by the wallet, it is not necessarily equal
+        // to TX_FEE
+        let lock_tx_bitcoin_fee = self
+            .bitcoin_wallet
+            .transaction_fee(lock_tx_id)
+            .await
+            .unwrap();
+
+        let btc_balance_after_swap = self.bitcoin_wallet.as_ref().balance().await.unwrap();
+        assert_eq!(
+            btc_balance_after_swap,
+            self.starting_balances.btc - self.swap_amounts.btc - lock_tx_bitcoin_fee
+        );
+
+        let xmr_balance_after_swap = self.monero_wallet.as_ref().get_balance().await.unwrap();
+        assert_eq!(xmr_balance_after_swap, self.starting_balances.xmr);
+    }
 }
 
 pub async fn test<T, F>(testfn: T)