diff --git a/roles/tests-integration/tests/common/sniffer.rs b/roles/tests-integration/tests/common/sniffer.rs index 985f1ba92f..aa53c46266 100644 --- a/roles/tests-integration/tests/common/sniffer.rs +++ b/roles/tests-integration/tests/common/sniffer.rs @@ -418,6 +418,21 @@ impl Sniffer { panic!("Impossible to accept dowsntream connection") } } + + /// used to block the test runtime + /// while we wait until Sniffer has a specific count of messages + pub async fn wait_for_message_count(&self, message_direction: MessageDirection, count: usize) { + loop { + let message_count = match message_direction { + MessageDirection::ToDownstream => self.messages_from_upstream.message_count(), + MessageDirection::ToUpstream => self.messages_from_downstream.message_count(), + }; + + if message_count >= count { + return; + } + } + } } // Utility macro to assert that the downstream and upstream roles have sent specific messages. @@ -589,6 +604,10 @@ impl MessagesAggregator { .unwrap() } + fn message_count(&self) -> usize { + self.messages.safe_lock(|messages| messages.len()).unwrap() + } + // The aggregator queues messages in FIFO order, so this function returns the oldest message in // the queue. //