ClickHouse · pamarcos · Jun 24, 2026 · Jun 24, 2026 · Jun 24, 2026
diff --git a/tests/integration/test_storage_rabbitmq/__init__.py b/tests/integration/test_storage_rabbitmq/__init__.py
@@ -0,0 +1,29 @@
+import logging
+import time
+
+import pytest
+
+
+def poll_direct_select_result(
+    instance, query, is_result_complete, timeout, ignore_error=True, sleep_time=0.05
+):
+    result = ""
+    previous_rows = 0
+    deadline = time.monotonic() + timeout
+    while time.monotonic() < deadline:
+        # Direct SELECT from RabbitMQ consumes rows, so accumulate all returned batches.
+        result += instance.query(query, ignore_error=ignore_error)
+        if is_result_complete(result):
+            return result
+
+        rows = len([line for line in result.splitlines() if line.strip()])
+        if rows > previous_rows:
+            deadline = time.monotonic() + timeout
+        previous_rows = rows
+        logging.debug(f"Collected rows: {rows}. Now {time.monotonic()}, deadline {deadline}")
+        time.sleep(sleep_time)
+
+    pytest.fail(
+        f"Time limit of {timeout} seconds reached without RabbitMQ consumption progress. "
+        "The result did not match the expected value."
+    )
diff --git a/tests/integration/test_storage_rabbitmq/test.py b/tests/integration/test_storage_rabbitmq/test.py
@@ -13,6 +13,7 @@
 from helpers.client import QueryRuntimeException
 from helpers.cluster import ClickHouseCluster
 from helpers.test_tools import TSV
+from . import poll_direct_select_result
 
 pytestmark = pytest.mark.timeout(1200)
 
@@ -63,6 +64,19 @@ def rabbitmq_check_result(result, check=False, reference=None):
         return TSV(result) == TSV(reference)
 
 
+def check_direct_select_result_polling(query, reference=None, timeout=DEFAULT_TIMEOUT_SEC):
+    if reference is None:
+        reference = "\n".join([f"{i}\t{i}" for i in range(50)])
+
+    result = poll_direct_select_result(
+        instance,
+        query,
+        lambda current_result: rabbitmq_check_result(current_result, reference=reference),
+        timeout,
+    )
+    rabbitmq_check_result(result, True, reference=reference)
+
+
 # Fixtures
 
 
@@ -175,21 +189,7 @@ def test_rabbitmq_select(rabbitmq_cluster, secure, db, unique):
     # The order of messages in select * from {db}.rabbitmq is not guaranteed, so sleep to collect everything in one select
     time.sleep(1)
 
-    result = ""
-    deadline = time.monotonic() + DEFAULT_TIMEOUT_SEC
-    while time.monotonic() < deadline:
-        result += instance.query(
-            f"SELECT * FROM {db}.rabbitmq ORDER BY key", ignore_error=True
-        )
-        if rabbitmq_check_result(result):
-            break
-        time.sleep(0.05)
-    else:
-        pytest.fail(
-            f"Time limit of {DEFAULT_TIMEOUT_SEC} seconds reached. The result did not match the expected value."
-        )
-
-    rabbitmq_check_result(result, True)
+    check_direct_select_result_polling(f"SELECT * FROM {db}.rabbitmq ORDER BY key")
 
 
 def test_rabbitmq_select_empty(rabbitmq_cluster, db, unique):
@@ -248,22 +248,7 @@ def test_rabbitmq_json_without_delimiter(rabbitmq_cluster, db, unique):
     connection.close()
     time.sleep(1)
 
-    result = ""
-    deadline = time.monotonic() + DEFAULT_TIMEOUT_SEC
-    while time.monotonic() < deadline:
-        result += instance.query(
-            f"SELECT * FROM {db}.rabbitmq ORDER BY key", ignore_error=True
-        )
-        if rabbitmq_check_result(result):
-            break
-
-        time.sleep(0.05)
-    else:
-        pytest.fail(
-            f"Time limit of {DEFAULT_TIMEOUT_SEC} seconds reached. The result did not match the expected value."
-        )
-
-    rabbitmq_check_result(result, True)
+    check_direct_select_result_polling(f"SELECT * FROM {db}.rabbitmq ORDER BY key")
 
 
 def test_rabbitmq_csv_with_delimiter(rabbitmq_cluster, db, unique):
@@ -298,22 +283,7 @@ def test_rabbitmq_csv_with_delimiter(rabbitmq_cluster, db, unique):
     connection.close()
     time.sleep(1)
 
-    result = ""
-    deadline = time.monotonic() + DEFAULT_TIMEOUT_SEC
-    while time.monotonic() < deadline:
-        result += instance.query(
-            f"SELECT * FROM {db}.rabbitmq ORDER BY key", ignore_error=True
-        )
-        if rabbitmq_check_result(result):
-            break
-
-        time.sleep(0.05)
-    else:
-        pytest.fail(
-            f"Time limit of {DEFAULT_TIMEOUT_SEC} seconds reached. The result did not match the expected value."
-        )
-
-    rabbitmq_check_result(result, True)
+    check_direct_select_result_polling(f"SELECT * FROM {db}.rabbitmq ORDER BY key")
 
 
 def test_rabbitmq_tsv_with_delimiter(rabbitmq_cluster, db, unique):
@@ -398,21 +368,7 @@ def test_rabbitmq_macros(rabbitmq_cluster, db, unique):
     connection.close()
     time.sleep(1)
 
-    result = ""
-    deadline = time.monotonic() + DEFAULT_TIMEOUT_SEC
-    while time.monotonic() < deadline:
-        result += instance.query(
-            f"SELECT * FROM {db}.rabbitmq ORDER BY key", ignore_error=True
-        )
-        if rabbitmq_check_result(result):
-            break
-        time.sleep(0.05)
-    else:
-        pytest.fail(
-            f"Time limit of {DEFAULT_TIMEOUT_SEC} seconds reached. The result did not match the expected value."
-        )
-
-    rabbitmq_check_result(result, True)
+    check_direct_select_result_polling(f"SELECT * FROM {db}.rabbitmq ORDER BY key")
 
 
 def test_rabbitmq_materialized_view(rabbitmq_cluster, db, unique):
@@ -585,6 +541,9 @@ def test_rabbitmq_many_materialized_views(rabbitmq_cluster, db, unique):
         channel.basic_publish(exchange=f"{unique}_mmv", routing_key="", body=message)
 
     is_check_passed = False
+    result1 = ""
+    result2 = ""
+    result3 = ""
     deadline = time.monotonic() + DEFAULT_TIMEOUT_SEC
     while time.monotonic() < deadline:
         result1 = instance.query(f"SELECT * FROM {db}.view1 ORDER BY key")

diff --git a/tests/integration/test_storage_rabbitmq/test_shutdown.py b/tests/integration/test_storage_rabbitmq/test_shutdown.py
@@ -3,6 +3,7 @@
 import logging
 import time
 from helpers.cluster import ClickHouseCluster
+from . import poll_direct_select_result
 
 cluster = ClickHouseCluster(__file__)
 
@@ -389,19 +390,20 @@ def test_rabbitmq_virtual_column_table(started_cluster):
         """
     )
 
-    result = ""
-    for _ in range(100):
-        result += instance.query(
-            """
-            SELECT key, value, _exchange_name, _table
-            FROM test_virt.rabbitmq_source
-            SETTINGS stream_like_engine_allow_direct_select=1
-            """
-        )
-        lines = [l for l in result.strip().split("\n") if l]
-        if len(lines) == 10:
-            break
-        time.sleep(0.5)
+    result = poll_direct_select_result(
+        instance,
+        """
+        SELECT key, value, _exchange_name, _table
+        FROM test_virt.rabbitmq_source
+        SETTINGS stream_like_engine_allow_direct_select=1
+        """,
+        lambda current_result: len(
+            [l for l in current_result.strip().split("\n") if l]
+        ) == 10,
+        timeout=50,
+        ignore_error=False,
+        sleep_time=0.5,
+    )
 
     lines = [l for l in result.strip().split("\n") if l]
     assert len(lines) == 10, f"Expected 10 rows, got {len(lines)}"