From afb908519770bd9057157853738c5e788093546b Mon Sep 17 00:00:00 2001
From: jbesraa <jbesraa@gmail.com>
Date: Fri, 10 Jan 2025 13:29:11 +0200
Subject: [PATCH] Use `wait_for_message_type_and_clean_queue`

in `header_timestamp_value_assertion_in_new_extended_mining_job` test.
---
 .../tests/pool_integration.rs                 | 29 +++++++------------
 1 file changed, 10 insertions(+), 19 deletions(-)

diff --git a/roles/tests-integration/tests/pool_integration.rs b/roles/tests-integration/tests/pool_integration.rs
index e6f3446d15..94c832a0dd 100644
--- a/roles/tests-integration/tests/pool_integration.rs
+++ b/roles/tests-integration/tests/pool_integration.rs
@@ -1,7 +1,10 @@
 use integration_tests_sv2::*;
 
 use crate::sniffer::MessageDirection;
-use const_sv2::{MESSAGE_TYPE_NEW_EXTENDED_MINING_JOB, MESSAGE_TYPE_NEW_TEMPLATE};
+use const_sv2::{
+    MESSAGE_TYPE_MINING_SET_NEW_PREV_HASH, MESSAGE_TYPE_NEW_EXTENDED_MINING_JOB,
+    MESSAGE_TYPE_NEW_TEMPLATE,
+};
 use roles_logic_sv2::{
     common_messages_sv2::{Protocol, SetupConnection},
     parsers::{AnyMessage, CommonMessages, Mining, PoolMessages, TemplateDistribution},
@@ -92,24 +95,12 @@ async fn header_timestamp_value_assertion_in_new_extended_mining_job() {
         }
         _ => panic!("SetNewPrevHash not found!"),
     };
-    // Assertions of messages between Pool and Translator Proxy (these are not necessary for the
-    // test itself, but they are used to pop from the sniffer's message queue)
-    assert_common_message!(
-        &pool_translator_sniffer.next_message_from_upstream(),
-        SetupConnectionSuccess
-    );
-    assert_mining_message!(
-        &pool_translator_sniffer.next_message_from_upstream(),
-        OpenExtendedMiningChannelSuccess
-    );
-    assert_mining_message!(
-        &pool_translator_sniffer.next_message_from_upstream(),
-        NewExtendedMiningJob
-    );
-    assert_mining_message!(
-        &pool_translator_sniffer.next_message_from_upstream(),
-        SetNewPrevHash
-    );
+    pool_translator_sniffer
+        .wait_for_message_type_and_clean_queue(
+            MessageDirection::ToDownstream,
+            MESSAGE_TYPE_MINING_SET_NEW_PREV_HASH,
+        )
+        .await;
     // Wait for a second NewExtendedMiningJob message
     pool_translator_sniffer
         .wait_for_message_type(