improve test method

2 files changed, 35 insertions, 5 deletions

diff --git a/ydb/core/testlib/common_helper.cpp b/ydb/core/testlib/common_helper.cpp
index 35bacbb6af..6b97f6d1e1 100644
--- a/ydb/core/testlib/common_helper.cpp
+++ b/ydb/core/testlib/common_helper.cpp
@@ -4,6 +4,7 @@
 #include <ydb/core/tx/schemeshard/schemeshard.h>
 
 #include <ydb/public/sdk/cpp/client/ydb_table/table.h>
+#include <ydb/public/lib/yson_value/ydb_yson_value.h>
 
 #include <library/cpp/testing/unittest/registar.h>
 
@@ -19,18 +20,20 @@ void THelper::WaitForSchemeOperation(TActorId sender, ui64 txId) {
     runtime.GrabEdgeEventRethrow<NSchemeShard::TEvSchemeShard::TEvNotifyTxCompletionResult>(sender);
 }
 
-void THelper::StartDataRequest(const TString& request, const bool expectSuccess) const {
+void THelper::StartDataRequest(const TString& request, const bool expectSuccess, TString* result) const {
     NYdb::NTable::TTableClient tClient(Server.GetDriver(),
         NYdb::NTable::TClientSettings().UseQueryCache(false).AuthToken("root@builtin"));
     auto expectation = expectSuccess;
     bool resultReady = false;
     bool* rrPtr = &resultReady;
-    tClient.CreateSession().Subscribe([rrPtr, request, expectation](NThreading::TFuture<NYdb::NTable::TCreateSessionResult> f) {
+    std::optional<NYdb::NTable::TDataQueryResult> resultInternal;
+    tClient.CreateSession().Subscribe([&resultInternal, rrPtr, request, expectation](NThreading::TFuture<NYdb::NTable::TCreateSessionResult> f) {
         auto session = f.GetValueSync().GetSession();
         session.ExecuteDataQuery(request
             , NYdb::NTable::TTxControl::BeginTx(NYdb::NTable::TTxSettings::SerializableRW()).CommitTx())
-            .Subscribe([rrPtr, expectation, request](NYdb::NTable::TAsyncDataQueryResult f)
+            .Subscribe([&resultInternal, rrPtr, expectation, request](NYdb::NTable::TAsyncDataQueryResult f)
                 {
+                    resultInternal = f.GetValueSync();
                     TStringStream ss;
                     f.GetValueSync().GetIssues().PrintTo(ss, false);
                     Cerr << "REQUEST=" << request << ";RESULT=" << ss.Str() << ";EXPECTATION=" << expectation << Endl;
@@ -39,11 +42,20 @@ void THelper::StartDataRequest(const TString& request, const bool expectSuccess)
                 });
         });
     const TInstant start = TInstant::Now();
-    while (!resultReady && start + TDuration::Seconds(200) > TInstant::Now()) {
+    while (!resultReady && start + TDuration::Seconds(60) > TInstant::Now()) {
         Server.GetRuntime()->SimulateSleep(TDuration::Seconds(1));
     }
     Cerr << "REQUEST=" << request << ";EXPECTATION=" << expectation << Endl;
     UNIT_ASSERT(resultReady);
+    UNIT_ASSERT(resultInternal);
+    if (result) {
+        TStringStream ss;
+        NYson::TYsonWriter writer(&ss, NYson::EYsonFormat::Text);
+        for (auto&& i : resultInternal->GetResultSets()) {
+            PrintResultSet(i, writer);
+        }
+        *result = ss.Str();
+    }
 }
 
 void THelper::StartSchemaRequest(const TString& request, const bool expectSuccess, const bool waiting) const {
@@ -96,4 +108,19 @@ void THelper::DropTable(const TString& tablePath) {
     runtime->DispatchEvents(options);
 }
 
+void THelper::PrintResultSet(const NYdb::TResultSet& resultSet, NYson::TYsonWriter& writer) const {
+    auto columns = resultSet.GetColumnsMeta();
+
+    NYdb::TResultSetParser parser(resultSet);
+    while (parser.TryNextRow()) {
+        writer.OnListItem();
+        writer.OnBeginList();
+        for (ui32 i = 0; i < columns.size(); ++i) {
+            writer.OnListItem();
+            FormatValueYson(parser.GetValue(i), writer);
+        }
+        writer.OnEndList();
+    }
+}
+
 }
diff --git a/ydb/core/testlib/common_helper.h b/ydb/core/testlib/common_helper.h
index d25393972d..0abc659b97 100644
--- a/ydb/core/testlib/common_helper.h
+++ b/ydb/core/testlib/common_helper.h
@@ -1,11 +1,14 @@
 #pragma once
 #include "test_client.h"
+#include <ydb/public/sdk/cpp/client/ydb_result/result.h>
+#include <library/cpp/yson/writer.h>
 
 namespace NKikimr::Tests::NCommon {
 
 class THelper {
 protected:
     void WaitForSchemeOperation(TActorId sender, ui64 txId);
+    void PrintResultSet(const NYdb::TResultSet& resultSet, NYson::TYsonWriter& writer) const;
 
     Tests::TServer& Server;
 public:
@@ -16,7 +19,7 @@ public:
 
     void DropTable(const TString& tablePath);
 
-    void StartDataRequest(const TString& request, const bool expectSuccess = true) const;
+    void StartDataRequest(const TString& request, const bool expectSuccess = true, TString* result = nullptr) const;
     void StartSchemaRequest(const TString& request, const bool expectSuccess = true, const bool waiting = true) const;
 };
 }