add waiting for resources snapshot (#1154)

1 files changed, 24 insertions, 12 deletions

diff --git a/ydb/core/kqp/ut/scan/kqp_scan_ut.cpp b/ydb/core/kqp/ut/scan/kqp_scan_ut.cpp
index 8309e4ea86..093889290b 100644
--- a/ydb/core/kqp/ut/scan/kqp_scan_ut.cpp
+++ b/ydb/core/kqp/ut/scan/kqp_scan_ut.cpp
@@ -2583,27 +2583,39 @@ Y_UNIT_TEST_SUITE(KqpRequestContext) {
         auto settings = TKikimrSettings()
             .SetAppConfig(AppCfg())
             .SetEnableScriptExecutionOperations(true)
-            .SetNodeCount(4);
+            .SetNodeCount(4)
+            .SetUseRealThreads(false);
         TKikimrRunner kikimr{settings};
         auto db = kikimr.GetTableClient();
-        
+
         NKikimr::NKqp::TKqpPlanner::UseMockEmptyPlanner = true;
         Y_DEFER {
             NKikimr::NKqp::TKqpPlanner::UseMockEmptyPlanner = false;  // just in case if test fails
         };
 
-        auto it = db.StreamExecuteScanQuery(R"(
-            SELECT Text, SUM(Key) AS Total FROM `/Root/EightShard`
-            GROUP BY Text
-            ORDER BY Total DESC;
-        )").GetValueSync();
+        {
+            TDispatchOptions opts;
+            opts.FinalEvents.emplace_back(NKikimr::NKqp::TKqpResourceInfoExchangerEvents::EvSendResources, 4);
+            kikimr.GetTestServer().GetRuntime()->DispatchEvents(opts);
+        }
+
+        auto it = kikimr.RunCall([&db] {
+            return db.StreamExecuteScanQuery(R"(
+                SELECT Text, SUM(Key) AS Total FROM `/Root/EightShard`
+                GROUP BY Text
+                ORDER BY Total DESC;
+            )").GetValueSync();
+        });
 
         UNIT_ASSERT(it.IsSuccess());
-        try {
-            auto yson = StreamResultToYson(it, true, NYdb::EStatus::PRECONDITION_FAILED, "TraceId");
-        } catch (const std::exception& ex) {
-            UNIT_ASSERT_C(false, "Exception NYdb::EStatus::PRECONDITION_FAILED not found or IssueMessage doesn't contain 'TraceId'");
-        }
+        kikimr.RunCall([&it] {
+            try {
+                auto yson = StreamResultToYson(it, true, NYdb::EStatus::PRECONDITION_FAILED, "TraceId");
+            } catch (const std::exception& ex) {
+                UNIT_ASSERT_C(false, "Exception NYdb::EStatus::PRECONDITION_FAILED not found or IssueMessage doesn't contain 'TraceId'");
+            }
+            return true;
+        });
 
         NKikimr::NKqp::TKqpPlanner::UseMockEmptyPlanner = false;
     }