diff --git a/roles/tests-integration/tests/common/sniffer.rs b/roles/tests-integration/tests/common/sniffer.rs index 985f1ba92..0e61fb4c9 100644 --- a/roles/tests-integration/tests/common/sniffer.rs +++ b/roles/tests-integration/tests/common/sniffer.rs @@ -418,6 +418,25 @@ 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 +608,12 @@ 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. //