Fix flaky domain tests

Signed-off-by: linning <[email protected]>
autonomys · Nov 10, 2023 · 8a3f52a · 8a3f52a
1 parent bab5695
commit 8a3f52a
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 26 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git 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.28"
-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
@@ -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, FutureExt, StreamExt};
+use futures::StreamExt;
 use jsonrpsee::RpcModule;
 use parking_lot::Mutex;
 use sc_block_builder::BlockBuilderProvider;
@@ -464,7 +464,10 @@ impl MockConsensusNode {
             self.acknowledgement_sender_subscribers
                 .retain(|subscriber| {
                     subscriber
+                        // Send 2 `acknowledgement_sender` to ensure all the domain worker's previous
+                        // tasks have finished
                         .unbounded_send(acknowledgement_sender.clone())
+                        .and_then(|_| subscriber.unbounded_send(acknowledgement_sender.clone()))
                         .is_ok()
                 });
             drop(acknowledgement_sender);
@@ -575,29 +578,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>> {