From c0aa06682cebf8e7942deda2480907f1b5d57d4c Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Thu, 6 Feb 2025 12:21:11 +1000
Subject: [PATCH] Add a timeout to test_domain_tx_propagate to diagnose hangs

---
 domains/client/domain-operator/src/tests.rs | 70 ++++++++++++---------
 1 file changed, 40 insertions(+), 30 deletions(-)

diff --git a/domains/client/domain-operator/src/tests.rs b/domains/client/domain-operator/src/tests.rs
index 72b93d564b..8d4781c163 100644
--- a/domains/client/domain-operator/src/tests.rs
+++ b/domains/client/domain-operator/src/tests.rs
@@ -72,6 +72,7 @@ use sp_transaction_pool::runtime_api::TaggedTransactionQueue;
 use sp_weights::Weight;
 use std::assert_matches::assert_matches;
 use std::collections::{BTreeMap, VecDeque};
+use std::future::IntoFuture;
 use std::sync::Arc;
 use std::time::Duration;
 use subspace_core_primitives::pot::PotOutput;
@@ -83,8 +84,23 @@ use subspace_test_service::{
 use tempfile::TempDir;
 use tracing::error;
 
+/// The general timeout for test operations that could hang.
 const TIMEOUT: Duration = Duration::from_mins(2);
 
+/// A trait that makes it easier to add a timeout to any future: `future.timeout().await?`.
+trait TestTimeout: IntoFuture {
+    fn timeout(self) -> tokio::time::Timeout<Self::IntoFuture>;
+}
+
+impl<F> TestTimeout for F
+where
+    F: IntoFuture,
+{
+    fn timeout(self) -> tokio::time::Timeout<Self::IntoFuture> {
+        tokio::time::timeout(TIMEOUT, self)
+    }
+}
+
 fn number_of(consensus_node: &MockConsensusNode, block_hash: Hash) -> u32 {
     consensus_node
         .client
@@ -1600,44 +1616,34 @@ async fn collected_receipts_should_be_on_the_same_branch_with_current_best_block
     );
 }
 
+// This test hangs occasionally, but we don't know which step hangs.
+// TODO: when the test is fixed, decide if we want to remove the timeouts.
 #[tokio::test(flavor = "multi_thread")]
-async fn test_domain_tx_propagate() {
-    let directory = TempDir::new().expect("Must be able to create temporary directory");
-
-    let mut builder = sc_cli::LoggerBuilder::new("");
-    builder.with_colors(false);
-    let _ = builder.init();
-
-    let tokio_handle = tokio::runtime::Handle::current();
-
-    // Start Ferdie
-    let mut ferdie = MockConsensusNode::run(
-        tokio_handle.clone(),
-        Ferdie,
-        BasePath::new(directory.path().join("ferdie")),
-    );
-
-    // Run Alice (a evm domain authority node)
-    let alice = domain_test_service::DomainNodeBuilder::new(
-        tokio_handle.clone(),
-        BasePath::new(directory.path().join("alice")),
-    )
-    .build_evm_node(Role::Authority, Alice, &mut ferdie)
-    .await;
+async fn test_domain_tx_propagate() -> Result<(), tokio::time::error::Elapsed> {
+    let (directory, mut ferdie, alice) = setup_evm_test_nodes(Ferdie).timeout().await?;
 
     // Run Bob (a evm domain full node)
     let mut bob = domain_test_service::DomainNodeBuilder::new(
-        tokio_handle,
+        tokio::runtime::Handle::current(),
         BasePath::new(directory.path().join("bob")),
     )
     .connect_to_domain_node(alice.addr.clone())
     .build_evm_node(Role::Full, Bob, &mut ferdie)
-    .await;
+    .timeout()
+    .await?;
+
+    produce_blocks!(ferdie, alice, 5, bob)
+        .timeout()
+        .await?
+        .unwrap();
 
-    produce_blocks!(ferdie, alice, 5, bob).await.unwrap();
-    while alice.sync_service.is_major_syncing() || bob.sync_service.is_major_syncing() {
-        tokio::time::sleep(std::time::Duration::from_millis(100)).await;
+    async {
+        while alice.sync_service.is_major_syncing() || bob.sync_service.is_major_syncing() {
+            tokio::time::sleep(std::time::Duration::from_millis(100)).await;
+        }
     }
+    .timeout()
+    .await?;
 
     let pre_bob_free_balance = alice.free_balance(bob.key.to_account_id());
     // Construct and send an extrinsic to bob, as bob is not a authority node, the extrinsic has
@@ -1646,7 +1652,8 @@ async fn test_domain_tx_propagate() {
         dest: Alice.to_account_id(),
         value: 123,
     })
-    .await
+    .timeout()
+    .await?
     .expect("Failed to send extrinsic");
 
     produce_blocks_until!(
@@ -1659,8 +1666,11 @@ async fn test_domain_tx_propagate() {
         },
         bob
     )
-    .await
+    .timeout()
+    .await?
     .unwrap();
+
+    Ok(())
 }
 
 #[tokio::test(flavor = "multi_thread")]