Retry rejected status in FlatQuery for tests stability

3 files changed, 20 insertions, 3 deletions

diff --git a/ydb/core/client/client_ut.cpp b/ydb/core/client/client_ut.cpp
index 849a20f874..3c4bf38832 100644
--- a/ydb/core/client/client_ut.cpp
+++ b/ydb/core/client/client_ut.cpp
@@ -1647,6 +1647,9 @@ Y_UNIT_TEST_SUITE(TClientTest) {
         WaitForFollowerStart(client, server.GetRuntime(), tabletId, TDuration::Seconds(5));
 
         TVector<ui32> followerNodes = client.GetFollowerNodes(server.GetRuntime(), tabletId);
+        while (followerNodes.empty()) {
+            followerNodes = client.GetFollowerNodes(server.GetRuntime(), tabletId);
+        }
         UNIT_ASSERT_VALUES_EQUAL(1, followerNodes.size());
         UNIT_ASSERT_VALUES_UNEQUAL(leaderNode, followerNodes[0]);
 
diff --git a/ydb/core/testlib/test_client.cpp b/ydb/core/testlib/test_client.cpp
index 0efb648e4b..987c7370f6 100644
--- a/ydb/core/testlib/test_client.cpp
+++ b/ydb/core/testlib/test_client.cpp
@@ -2089,7 +2089,21 @@ namespace Tests {
 
     bool TClient::FlatQuery(const TString &query, TFlatQueryOptions& opts, NKikimrMiniKQL::TResult &result, const NKikimrClient::TResponse& expectedResponse) {
         NKikimrClient::TResponse response;
-        FlatQueryRaw(query, opts, response);
+        if (expectedResponse.HasStatus() && expectedResponse.GetStatus() == NMsgBusProxy::MSTATUS_OK) {
+            // Client is expecting OK, retry REJECTED replies during restarts and splits
+            for (int i = 0; i < 5; ++i) {
+                if (i != 0) {
+                    response.Clear();
+                    Cerr << "Retrying rejected query..." << Endl;
+                }
+                FlatQueryRaw(query, opts, response);
+                if (response.GetStatus() != NMsgBusProxy::MSTATUS_REJECTED) {
+                    break;
+                }
+            }
+        } else {
+            FlatQueryRaw(query, opts, response);
+        }
 
         if (!response.GetDataShardErrors().empty()) {
             Cerr << "DataShardErrors:" << Endl << response.GetDataShardErrors() << Endl;
@@ -2347,7 +2361,7 @@ namespace Tests {
 
         TVector<ui32> followerNodes;
         for (const NKikimrHive::TTabletInfo& tablet : res.GetTablets()) {
-            if (tablet.GetTabletID() == tabletId && tablet.HasFollowerID()) {
+            if (tablet.GetTabletID() == tabletId && tablet.HasFollowerID() && tablet.GetNodeID() != 0) {
                 followerNodes.push_back(NodeIdToIndex(runtime, tablet.GetNodeID()));
             }
         }
diff --git a/ydb/core/testlib/test_client.h b/ydb/core/testlib/test_client.h
index 4bf212ca1d..21de04466e 100644
--- a/ydb/core/testlib/test_client.h
+++ b/ydb/core/testlib/test_client.h
@@ -525,7 +525,7 @@ namespace Tests {
 
         ui32 NodeIdToIndex(TTestActorRuntime* runtime, ui32 id) {
             ui32 offset = runtime->GetNodeId(0);
-            Y_VERIFY(id >= offset);
+            Y_VERIFY(id >= offset, "NodeId# %" PRIu32 " offset# %" PRIu32, id, offset);
             return id - offset;
         }