From 2ec0a9cfeaf850c31f1567da1884ad77fae0bcdd Mon Sep 17 00:00:00 2001
From: Vitaly Baranov <vitlibar@clickhouse.com>
Date: Wed, 3 Jul 2024 19:25:44 +0200
Subject: [PATCH] Fix test test_grpc_protocol/test.py::test_progress

---
 tests/integration/test_grpc_protocol/test.py | 60 ++++++++------------
 1 file changed, 23 insertions(+), 37 deletions(-)

diff --git a/tests/integration/test_grpc_protocol/test.py b/tests/integration/test_grpc_protocol/test.py
index 851da99acf3..1e4ae7f0288 100644
--- a/tests/integration/test_grpc_protocol/test.py
+++ b/tests/integration/test_grpc_protocol/test.py
@@ -369,47 +369,33 @@ def test_progress():
         "SELECT number, sleep(0.31) FROM numbers(8) SETTINGS max_block_size=2, interactive_delay=100000",
         stream_output=True,
     )
-    results = list(results)
-    for result in results:
-        result.time_zone = ""
-        result.query_id = ""
-    # print(results)
 
-    # Note: We can't convert those messages to string like `results = str(results)` and then compare it as a string
-    # because str() can serialize a protobuf message with any order of fields.
-    expected_results = [
-        clickhouse_grpc_pb2.Result(
-            output_format="TabSeparated",
-            progress=clickhouse_grpc_pb2.Progress(
-                read_rows=2, read_bytes=16, total_rows_to_read=8
-            ),
-        ),
-        clickhouse_grpc_pb2.Result(output=b"0\t0\n1\t0\n"),
-        clickhouse_grpc_pb2.Result(
-            progress=clickhouse_grpc_pb2.Progress(read_rows=2, read_bytes=16)
-        ),
-        clickhouse_grpc_pb2.Result(output=b"2\t0\n3\t0\n"),
-        clickhouse_grpc_pb2.Result(
-            progress=clickhouse_grpc_pb2.Progress(read_rows=2, read_bytes=16)
-        ),
-        clickhouse_grpc_pb2.Result(output=b"4\t0\n5\t0\n"),
-        clickhouse_grpc_pb2.Result(
-            progress=clickhouse_grpc_pb2.Progress(read_rows=2, read_bytes=16)
-        ),
-        clickhouse_grpc_pb2.Result(output=b"6\t0\n7\t0\n"),
-        clickhouse_grpc_pb2.Result(
-            stats=clickhouse_grpc_pb2.Stats(
-                rows=8,
-                blocks=4,
-                allocated_bytes=1092,
-            )
-        ),
+    # Note: We can't compare results using a statement like `assert results == expected_results`
+    # because `results` can come in slightly different order.
+    # So we compare `outputs` and `progresses` separately and not `results` as a whole.
+
+    outputs = [i.output for i in results if i.output]
+    progresses = [i.progress for i in results if i.HasField("progress")]
+
+    # print(outputs)
+    # print(progresses)
+
+    expected_outputs = [
+        b"0\t0\n1\t0\n",
+        b"2\t0\n3\t0\n",
+        b"4\t0\n5\t0\n",
+        b"6\t0\n7\t0\n",
     ]
 
-    # Stats data can be returned, which broke the test
-    results = [i for i in results if not isinstance(i, clickhouse_grpc_pb2.Stats)]
+    expected_progresses = [
+        clickhouse_grpc_pb2.Progress(read_rows=2, read_bytes=16, total_rows_to_read=8),
+        clickhouse_grpc_pb2.Progress(read_rows=2, read_bytes=16),
+        clickhouse_grpc_pb2.Progress(read_rows=2, read_bytes=16),
+        clickhouse_grpc_pb2.Progress(read_rows=2, read_bytes=16),
+    ]
 
-    assert results == expected_results
+    assert outputs == expected_outputs
+    assert progresses == expected_progresses
 
 
 def test_session_settings():