From e366d3a73b7cf33a3216e7123f055bb14b7558d4 Mon Sep 17 00:00:00 2001
From: Franck Royer <franck@coblox.tech>
Date: Mon, 21 Dec 2020 11:21:53 +1100
Subject: [PATCH] Join all futures to avoid hanging tests

---
 swap/tests/happy_path.rs               | 18 ++++++++++-------
 swap/tests/happy_path_restart_alice.rs | 10 +++++-----
 swap/tests/happy_path_restart_bob.rs   |  6 ------
 swap/tests/punish.rs                   | 27 ++++++++++++++++++++------
 4 files changed, 37 insertions(+), 24 deletions(-)

diff --git a/swap/tests/happy_path.rs b/swap/tests/happy_path.rs
index e4571f5a..8531e1c1 100644
--- a/swap/tests/happy_path.rs
+++ b/swap/tests/happy_path.rs
@@ -1,5 +1,8 @@
 use crate::testutils::{init_alice, init_bob};
-use futures::future::try_join;
+use futures::{
+    future::{join, select},
+    FutureExt,
+};
 use get_port::get_port;
 use libp2p::Multiaddr;
 use rand::rngs::OsRng;
@@ -82,11 +85,11 @@ async fn happy_path() {
         config,
         Uuid::new_v4(),
         alice_db,
-    );
+    )
+    .boxed();
 
     let alice_peer_id = alice_event_loop.peer_id();
-
-    let _alice_swarm_fut = tokio::spawn(async move { alice_event_loop.run().await });
+    let alice_fut = select(alice_swap_fut, alice_event_loop.run().boxed());
 
     let bob_swap_fut = bob::swap::swap(
         bob_state,
@@ -98,11 +101,12 @@ async fn happy_path() {
         Uuid::new_v4(),
         alice_peer_id,
         alice_multiaddr,
-    );
+    )
+    .boxed();
 
-    let _bob_swarm_fut = tokio::spawn(async move { bob_event_loop.run().await });
+    let bob_fut = select(bob_swap_fut, bob_event_loop.run().boxed());
 
-    try_join(alice_swap_fut, bob_swap_fut).await.unwrap();
+    join(alice_fut, bob_fut).await;
 
     let btc_alice_final = alice_btc_wallet.as_ref().balance().await.unwrap();
     let btc_bob_final = bob_btc_wallet.as_ref().balance().await.unwrap();
diff --git a/swap/tests/happy_path_restart_alice.rs b/swap/tests/happy_path_restart_alice.rs
index e32035ce..4d277509 100644
--- a/swap/tests/happy_path_restart_alice.rs
+++ b/swap/tests/happy_path_restart_alice.rs
@@ -76,7 +76,7 @@ async fn given_alice_restarts_after_encsig_is_learned_resume_swap() {
     let bob_btc_wallet_clone = bob_btc_wallet.clone();
     let bob_xmr_wallet_clone = bob_xmr_wallet.clone();
 
-    let _ = tokio::spawn(bob::swap::swap(
+    let bob_fut = bob::swap::swap(
         bob_state,
         bob_event_loop_handle,
         bob_db,
@@ -86,14 +86,14 @@ async fn given_alice_restarts_after_encsig_is_learned_resume_swap() {
         Uuid::new_v4(),
         alice_peer_id,
         alice_multiaddr.clone(),
-    ));
-
-    let _bob_swarm_fut = tokio::spawn(async move { bob_event_loop.run().await });
+    );
 
     let alice_db_datadir = tempdir().unwrap();
     let alice_db = Database::open(alice_db_datadir.path()).unwrap();
 
-    let _alice_swarm_fut = tokio::spawn(async move { alice_event_loop.run().await });
+    tokio::spawn(async move { alice_event_loop.run().await });
+    tokio::spawn(bob_fut);
+    tokio::spawn(bob_event_loop.run());
 
     let alice_swap_id = Uuid::new_v4();
 
diff --git a/swap/tests/happy_path_restart_bob.rs b/swap/tests/happy_path_restart_bob.rs
index 69f1b079..44ad1cb6 100644
--- a/swap/tests/happy_path_restart_bob.rs
+++ b/swap/tests/happy_path_restart_bob.rs
@@ -261,12 +261,6 @@ async fn given_bob_restarts_after_xmr_is_locked_resume_swap() {
         _ = bob_event_loop_1.run() => panic!("The event loop should never finish")
     };
 
-    // let tx_lock_id = if let AliceState::BtcRefunded { state3, .. } = alice_state
-    // {     state3.tx_lock.txid()
-    // } else {
-    //     panic!(format!("Alice in unexpected state: {}", alice_state));
-    // };
-
     let (bob_event_loop_2, bob_event_loop_handle_2) = testutils::init_bob_event_loop();
 
     let bob_fut = bob::swap::swap(
diff --git a/swap/tests/punish.rs b/swap/tests/punish.rs
index 91eb3164..0f8afaa4 100644
--- a/swap/tests/punish.rs
+++ b/swap/tests/punish.rs
@@ -1,5 +1,8 @@
 use crate::testutils::{init_alice, init_bob};
-use futures::future::try_join;
+use futures::{
+    future::{join, select, Either},
+    FutureExt,
+};
 use get_port::get_port;
 use libp2p::Multiaddr;
 use rand::rngs::OsRng;
@@ -81,9 +84,10 @@ async fn alice_punishes_if_bob_never_acts_after_fund() {
         bob_xmr_wallet.clone(),
         OsRng,
         Uuid::new_v4(),
-    );
+    )
+    .boxed();
 
-    let _bob_swarm_fut = tokio::spawn(async move { bob_event_loop.run().await });
+    let bob_fut = select(bob_btc_locked_fut, bob_event_loop.run().boxed());
 
     let alice_fut = alice::swap::swap(
         alice_state,
@@ -93,12 +97,23 @@ async fn alice_punishes_if_bob_never_acts_after_fund() {
         Config::regtest(),
         Uuid::new_v4(),
         alice_db,
-    );
+    )
+    .boxed();
 
-    let _alice_swarm_fut = tokio::spawn(async move { alice_event_loop.run().await });
+    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) = try_join(alice_fut, bob_btc_locked_fut).await.unwrap();
+    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::Punished));
     let bob_state3 = if let BobState::BtcLocked(state3, ..) = bob_state {