ydb_persqueue_public: fix flapping test (#4723)

Author: qyryq

ydb/public/sdk/cpp/client/ydb_persqueue_public/ut/retry_policy_ut.cpp

@@ -379,11 +379,11 @@ Y_UNIT_TEST_SUITE(RetryPolicy) {
     }
     Y_UNIT_TEST(RetryWithBatching) {
         auto setup = std::make_shared<TPersQueueYdbSdkTestSetup>(TEST_CASE_NAME);
-        auto settings = setup->GetWriteSessionSettings();
         auto retryPolicy = std::make_shared<TYdbPqTestRetryPolicy>();
-        settings.BatchFlushInterval(TDuration::Seconds(1000)); // Batch on size, not on time.
-        settings.BatchFlushSizeBytes(100);
-        settings.RetryPolicy(retryPolicy);
+        auto settings = setup->GetWriteSessionSettings()
+            .BatchFlushInterval(TDuration::Seconds(1000)) // Batch on size, not on time.
+            .BatchFlushSizeBytes(100)
+            .RetryPolicy(retryPolicy);
         auto& client = setup->GetPersQueueClient();
         auto writer = client.CreateWriteSession(settings);
         auto event = *writer->GetEvent(true);
@@ -393,6 +393,8 @@ Y_UNIT_TEST_SUITE(RetryPolicy) {
         TString message = "1234567890";
         ui64 seqNo = 0;
         setup->KickTablets();
+        setup->WaitForTabletsDown();
+
         writer->Write(std::move(continueToken), message, ++seqNo);
         retryPolicy->ExpectBreakDown();
         retryPolicy->WaitForRetriesSync(3);

ydb/public/sdk/cpp/client/ydb_persqueue_public/ut/ut_utils/sdk_test_setup.h

@@ -224,7 +224,18 @@ class SDKTestSetup {
             Server.AnnoyingClient->KickNodeInHive(Server.CleverServer->GetRuntime(), i);
         }
     }
-   
+
+    void WaitForTabletsDown() {
+        // After calling KickTablets wait until the tablets are in fact dead.
+
+        auto describeResult = Server.AnnoyingClient->Ls(GetTestTopicPath());
+        UNIT_ASSERT_C(describeResult->Record.GetPathDescription().HasPersQueueGroup(), describeResult->Record);
+        auto persQueueGroup = describeResult->Record.GetPathDescription().GetPersQueueGroup();
+        for (const auto& p : persQueueGroup.GetPartitions()) {
+            Server.AnnoyingClient->WaitForTabletDown(Server.CleverServer->GetRuntime(), p.GetTabletId(), true, TDuration::Max());
+        }
+    }
+
     void AllowTablets() {
         for (ui32 i = 0; i < Server.CleverServer->StaticNodes() + Server.CleverServer->DynamicNodes(); i++) {
             Server.AnnoyingClient->MarkNodeInHive(Server.CleverServer->GetRuntime(), i, true);

ydb/public/sdk/cpp/client/ydb_persqueue_public/ut/ut_utils/test_server.h

@@ -137,12 +137,12 @@ class TTestServer {
         if (killPq)
         {
             THashSet<ui64> restartedTablets;
-            for (const auto& p : persQueueGroup.GetPartitions())
-                if (restartedTablets.insert(p.GetTabletId()).second)
-                {
+            for (const auto& p : persQueueGroup.GetPartitions()) {
+                if (restartedTablets.insert(p.GetTabletId()).second) {
                     Log << TLOG_INFO << "Kill PQ tablet " << p.GetTabletId();
                     AnnoyingClient->KillTablet(*CleverServer, p.GetTabletId());
                 }
+            }
         }
 
         CleverServer->GetRuntime()->DispatchEvents();