From d508bc2e5c0b543a5eaff6f3a2791274f40974d5 Mon Sep 17 00:00:00 2001
From: linning <linningde25@gmail.com>
Date: Sat, 6 Jan 2024 02:47:19 +0800
Subject: [PATCH] Fix flaky domain tests

Signed-off-by: linning <linningde25@gmail.com>
---
 Cargo.lock                            |  1 -
 test/subspace-test-service/Cargo.toml |  1 -
 test/subspace-test-service/src/lib.rs | 36 +++++++--------------------
 3 files changed, 9 insertions(+), 29 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 0a4aea3c00..5c60ddc6df 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -12269,7 +12269,6 @@ dependencies = [
  "cross-domain-message-gossip",
  "domain-runtime-primitives",
  "futures",
- "futures-timer",
  "jsonrpsee",
  "pallet-domains",
  "parity-scale-codec",
diff --git a/test/subspace-test-service/Cargo.toml b/test/subspace-test-service/Cargo.toml
index e171eca1b8..bf320057ad 100644
--- a/test/subspace-test-service/Cargo.toml
+++ b/test/subspace-test-service/Cargo.toml
@@ -20,7 +20,6 @@ cross-domain-message-gossip = { version = "0.1.0", path = "../../domains/client/
 codec = { package = "parity-scale-codec", version = "3.2.1", features = ["derive"] }
 domain-runtime-primitives = { version = "0.1.0", path = "../../domains/primitives/runtime" }
 futures = "0.3.29"
-futures-timer = "3.0.1"
 jsonrpsee = { version = "0.16.3", features = ["server"] }
 rand = "0.8.5"
 pallet-domains = { version = "0.1.0", path = "../../crates/pallet-domains" }
diff --git a/test/subspace-test-service/src/lib.rs b/test/subspace-test-service/src/lib.rs
index 0b4c4d4d34..1ad5a7e9a8 100644
--- a/test/subspace-test-service/src/lib.rs
+++ b/test/subspace-test-service/src/lib.rs
@@ -22,7 +22,7 @@ use codec::{Decode, Encode};
 use cross_domain_message_gossip::GossipWorkerBuilder;
 use domain_runtime_primitives::opaque::{Block as DomainBlock, Header as DomainHeader};
 use futures::channel::mpsc;
-use futures::{select, Future, FutureExt, StreamExt};
+use futures::{Future, StreamExt};
 use jsonrpsee::RpcModule;
 use parking_lot::Mutex;
 use sc_block_builder::BlockBuilderProvider;
@@ -241,7 +241,7 @@ pub struct MockConsensusNode {
     next_slot: u64,
     /// The slot notification subscribers
     #[allow(clippy::type_complexity)]
-    new_slot_notification_subscribers: Vec<TracingUnboundedSender<(Slot, Randomness)>>,
+    new_slot_notification_subscribers: Vec<mpsc::UnboundedSender<(Slot, Randomness)>>,
     /// The acknowledgement sender subscribers
     #[allow(clippy::type_complexity)]
     acknowledgement_sender_subscribers: Vec<TracingUnboundedSender<mpsc::Sender<()>>>,
@@ -441,8 +441,8 @@ impl MockConsensusNode {
     }
 
     /// Subscribe the new slot notification
-    pub fn new_slot_notification_stream(&mut self) -> TracingUnboundedReceiver<(Slot, Randomness)> {
-        let (tx, rx) = tracing_unbounded("subspace_new_slot_notification_stream", 100);
+    pub fn new_slot_notification_stream(&mut self) -> mpsc::UnboundedReceiver<(Slot, Randomness)> {
+        let (tx, rx) = mpsc::unbounded();
         self.new_slot_notification_subscribers.push(tx);
         rx
     }
@@ -624,29 +624,11 @@ impl MockConsensusNode {
         &self,
         parent_number: NumberFor<Block>,
     ) -> Vec<<Block as BlockT>::Extrinsic> {
-        let mut t1 = self.transaction_pool.ready_at(parent_number).fuse();
-        let mut t2 = futures_timer::Delay::new(time::Duration::from_micros(100)).fuse();
-        let pending_iterator = select! {
-            res = t1 => res,
-            _ = t2 => {
-                tracing::warn!(
-                    "Timeout fired waiting for transaction pool at #{}, proceeding with production.",
-                    parent_number,
-                );
-                self.transaction_pool.ready()
-            }
-        };
-        let pushing_duration = time::Duration::from_micros(500);
-        let start = time::Instant::now();
-        let mut extrinsics = Vec::new();
-        for pending_tx in pending_iterator {
-            if start.elapsed() >= pushing_duration {
-                break;
-            }
-            let pending_tx_data = pending_tx.data().clone();
-            extrinsics.push(pending_tx_data);
-        }
-        extrinsics
+        self.transaction_pool
+            .ready_at(parent_number)
+            .await
+            .map(|pending_tx| pending_tx.data().clone())
+            .collect()
     }
 
     async fn mock_inherent_data(slot: Slot) -> Result<InherentData, Box<dyn Error>> {