Unmute watermark tests (#29121)

3 files changed, 20 insertions, 18 deletions

diff --git a/.github/config/muted_ya.txt b/.github/config/muted_ya.txt
index 8397f4cec2d..ba68db75371 100644
--- a/.github/config/muted_ya.txt
+++ b/.github/config/muted_ya.txt
@@ -11,7 +11,6 @@ ydb/core/blobstorage/ut_blobstorage/ut_huge unittest.[*/*] chunk
 ydb/core/client/ut TObjectStorageListingTest.TestSkipShards
 ydb/core/cms/ut_sentinel_unstable TSentinelUnstableTests.BSControllerCantChangeStatus
 ydb/core/cms/ut_sentinel_unstable unittest.sole chunk
-ydb/core/fq/libs/row_dispatcher/format_handler/ut TestFormatHandler.ManyRawClients
 ydb/core/kqp/ut/close_with_load KqpService.CloseSessionsWithLoad
 ydb/core/kqp/ut/close_with_load unittest.[*/*] chunk
 ydb/core/kqp/ut/federated_query/datastreams KqpFederatedQueryDatastreams.CheckpointsPropagationWithGroupByHop
@@ -179,7 +178,6 @@ ydb/tests/compatibility/streaming test_streaming.py.TestStreamingRollingUpgradeA
 ydb/tests/fq/mem_alloc test_scheduling.py.TestSchedule.test_skip_busy[kikimr0]
 ydb/tests/fq/multi_plane py3test.[test_retry.py] chunk
 ydb/tests/fq/multi_plane test_retry.py.TestRetry.test_low_rate[kikimr0]
-ydb/tests/fq/pq_async_io/ut TDqPqReadActorTest.TestSaveLoadPqRead
 ydb/tests/fq/s3 test_bindings_0.py.TestBindings.test_modify_connection_with_a_lot_of_bindings[v2-kikimr_settings0-client0]
 ydb/tests/fq/streaming py3test.sole chunk
 ydb/tests/fq/streaming test_streaming.py.TestStreamingInYdb.test_json_errors
diff --git a/ydb/core/fq/libs/row_dispatcher/format_handler/ut/format_handler_ut.cpp b/ydb/core/fq/libs/row_dispatcher/format_handler/ut/format_handler_ut.cpp
index d6a77c7b21c..6383597ce0f 100644
--- a/ydb/core/fq/libs/row_dispatcher/format_handler/ut/format_handler_ut.cpp
+++ b/ydb/core/fq/libs/row_dispatcher/format_handler/ut/format_handler_ut.cpp
@@ -195,7 +195,11 @@ public:
         Clients.emplace_back(client);
 
         if (!client->IsStarted()) {
-            const auto ev = Runtime.GrabEdgeEvent<NActors::TEvents::TEvPing>(CompileNotifier, TDuration::Seconds(5));
+            const auto timeout = TDuration::Seconds(5);
+
+            // wait for purecalc compile response to arrive before TEvPing
+            Sleep(timeout);
+            const auto ev = Runtime.GrabEdgeEvent<NActors::TEvents::TEvPing>(CompileNotifier, timeout);
             UNIT_ASSERT_C(ev, "Compilation is not performed for purecalc program");
         }
 
@@ -239,23 +243,20 @@ public:
 
     TCallback BatchCheck(TVector<TMessages> messages) const {
         return [this, expectedIndex = 0ull, expectedMessages = std::move(messages)](NActors::TActorId clientId, TQueue<TDataBatch>&& data) mutable {
-            UNIT_ASSERT_VALUES_EQUAL(data.size(), 1);
-            auto [actualMessages, actualOffsets, actualWatermark] = data.front();
-            data.pop();
+            while (!data.empty()) {
+                auto [actualMessages, actualOffsets, actualWatermark] = data.front();
+                data.pop();
 
-            UNIT_ASSERT_LT_C(expectedIndex, expectedMessages.size(), "Expected less messages");
-            auto [expectedOffsets, expectedWatermark, expectedBatch] = expectedMessages[expectedIndex++];
+                UNIT_ASSERT_LT_C(expectedIndex, expectedMessages.size(), "Expected less messages, clientId: " << clientId << ", got " << data.size() << " batches");
+                auto [expectedOffsets, expectedWatermark, expectedBatch] = expectedMessages[expectedIndex++];
 
-            UNIT_ASSERT_VALUES_EQUAL_C(expectedOffsets.size(), actualOffsets.size(), "clientId: " << clientId << ", expectedIndex: " << expectedIndex - 1);
-            size_t i = 0;
-            for (auto actualOffset : actualOffsets) {
-                UNIT_ASSERT_VALUES_EQUAL_C(expectedOffsets[i], actualOffset, "clientId: " << clientId << ", expectedIndex: " << expectedIndex - 1 << ", i: " << i);
-                ++i;
-            }
-            UNIT_ASSERT_VALUES_EQUAL_C(expectedWatermark, actualWatermark, "clientId: " << clientId << ", expectedIndex: " << expectedIndex - 1);
+                UNIT_ASSERT_VALUES_EQUAL_C(expectedOffsets, actualOffsets, "clientId: " << clientId << ", expectedIndex: " << expectedIndex - 1);
+
+                UNIT_ASSERT_VALUES_EQUAL_C(expectedWatermark, actualWatermark, "clientId: " << clientId << ", expectedIndex: " << expectedIndex - 1);
 
-            if (!expectedBatch.Rows.empty()) {
-                CheckMessageBatch(actualMessages, expectedBatch);
+                if (!expectedBatch.Rows.empty()) {
+                    CheckMessageBatch(actualMessages, expectedBatch);
+                }
             }
         };
     }
diff --git a/ydb/tests/fq/pq_async_io/ut_helpers.h b/ydb/tests/fq/pq_async_io/ut_helpers.h
index 7a872ee4c6e..6af6e270189 100644
--- a/ydb/tests/fq/pq_async_io/ut_helpers.h
+++ b/ydb/tests/fq/pq_async_io/ut_helpers.h
@@ -68,7 +68,10 @@ struct TPqIoTestFixture : public NUnitTest::TBaseFixture {
     }
 
     void LoadSource(const TSourceState& state) const {
-        return CaSetup->LoadSource(state);
+        CaSetup->LoadSource(state);
+
+        // Wait for reader to reconnect and resume from checkpoint
+        Sleep(TDuration::Seconds(10));
     }