Extract CountRows() to generic ut helpers (#19009)

6 files changed, 27 insertions, 37 deletions

diff --git a/ydb/core/tx/schemeshard/ut_helpers/helpers.cpp b/ydb/core/tx/schemeshard/ut_helpers/helpers.cpp
index 22feb17243a..d378c568650 100644
--- a/ydb/core/tx/schemeshard/ut_helpers/helpers.cpp
+++ b/ydb/core/tx/schemeshard/ut_helpers/helpers.cpp
@@ -18,6 +18,7 @@
 #include <ydb/public/api/protos/ydb_export.pb.h>
 #include <ydb/core/protos/schemeshard/operations.pb.h>
 #include <ydb/core/protos/auth.pb.h>
+#include <ydb/public/lib/deprecated/kicli/kicli.h>
 #include <ydb/public/sdk/cpp/include/ydb-cpp-sdk/client/table/table.h>
 
 #include <library/cpp/testing/unittest/registar.h>
@@ -2688,6 +2689,24 @@ namespace NSchemeShardUT_Private {
         return result;
     }
 
+    ui32 CountRows(TTestActorRuntime& runtime, ui64 schemeshardId, const TString& table) {
+        auto tableDesc = DescribePath(runtime, schemeshardId, table, true, false, true);
+        const auto& pathDesc = tableDesc.GetPathDescription();
+        const auto& key = pathDesc.GetTable().GetKeyColumnNames();
+        ui32 rows = 0;
+        for (const auto& x : pathDesc.GetTablePartitions()) {
+            auto result = ReadTable(runtime, x.GetDatashardId(), pathDesc.GetSelf().GetName(),
+                {key.begin(), key.end()}, {pathDesc.GetTable().GetKeyColumnNames()[0]});
+            auto value = NClient::TValue::Create(result);
+            rows += value["Result"]["List"].Size();
+        }
+        return rows;
+    }
+
+    ui32 CountRows(TTestActorRuntime& runtime, const TString& table) {
+        return CountRows(runtime, TTestTxConfig::SchemeShard, table);
+    }
+
     void WriteVectorTableRows(TTestActorRuntime& runtime, ui64 schemeShardId, ui64 txId, const TString & tablePath, bool withValue, ui32 shard, ui32 min, ui32 max) {
         TVector<TCell> cells;
         ui8 str[6] = { 0 };
diff --git a/ydb/core/tx/schemeshard/ut_helpers/helpers.h b/ydb/core/tx/schemeshard/ut_helpers/helpers.h
index aa09634b5d2..25863aece6c 100644
--- a/ydb/core/tx/schemeshard/ut_helpers/helpers.h
+++ b/ydb/core/tx/schemeshard/ut_helpers/helpers.h
@@ -658,6 +658,9 @@ namespace NSchemeShardUT_Private {
     NKikimrMiniKQL::TResult ReadTable(TTestActorRuntime& runtime, ui64 tabletId,
         const TString& table, const TVector<TString>& pk, const TVector<TString>& columns, const TString& rangeFlags = "");
 
+    ui32 CountRows(TTestActorRuntime& runtime, ui64 schemeshardId, const TString& table);
+    ui32 CountRows(TTestActorRuntime& runtime, const TString& table);
+
     void WriteVectorTableRows(TTestActorRuntime& runtime, ui64 schemeShardId, ui64 txId, const TString & tablePath,
         bool withValue, ui32 shard, ui32 min, ui32 max);
 
diff --git a/ydb/core/tx/schemeshard/ut_helpers/ya.make b/ydb/core/tx/schemeshard/ut_helpers/ya.make
index eabe089a44c..fb1af4a0088 100644
--- a/ydb/core/tx/schemeshard/ut_helpers/ya.make
+++ b/ydb/core/tx/schemeshard/ut_helpers/ya.make
@@ -21,6 +21,7 @@ PEERDIR(
     ydb/core/tx/tx_proxy
     ydb/public/lib/scheme_types
     yql/essentials/public/issue
+    ydb/public/lib/deprecated/kicli
     ydb/public/sdk/cpp/src/client/driver
     ydb/public/sdk/cpp/src/client/table
 )
diff --git a/ydb/core/tx/schemeshard/ut_index/ut_async_index.cpp b/ydb/core/tx/schemeshard/ut_index/ut_async_index.cpp
index f6c737b6212..707052da387 100644
--- a/ydb/core/tx/schemeshard/ut_index/ut_async_index.cpp
+++ b/ydb/core/tx/schemeshard/ut_index/ut_async_index.cpp
@@ -4,7 +4,6 @@
 #include <ydb/core/tx/schemeshard/ut_helpers/helpers.h>
 #include <ydb/core/tx/schemeshard/ut_helpers/test_with_reboots.h>
 #include <ydb/core/testlib/tablet_helpers.h>
-#include <ydb/public/lib/deprecated/kicli/kicli.h>
 #include <ydb/public/sdk/cpp/include/ydb-cpp-sdk/client/table/table.h>
 
 using namespace NKikimr;
@@ -126,37 +125,15 @@ Y_UNIT_TEST_SUITE(TAsyncIndexTests) {
         ui32 ExpectedRecords;
     };
 
-    template <typename C>
-    ui32 CountRows(TTestActorRuntime& runtime, const TTableTraits& table, const C& partitions) {
-        ui32 rows = 0;
-
-        for (const auto& x : partitions) {
-            auto result = ReadTable(runtime, x.GetDatashardId(), SplitPath(table.Path).back(), table.Key, table.Columns);
-            auto value = NClient::TValue::Create(result);
-            rows += value["Result"]["List"].Size();
-        }
-
-        return rows;
-    }
-
     bool CheckWrittenToIndex(TTestActorRuntime& runtime, const TTableTraits& mainTable, const TTableTraits& indexTable) {
-        bool writtenToMainTable = false;
-        {
-            auto tableDesc = DescribePath(runtime, mainTable.Path, true, true);
-            const auto& tablePartitions = tableDesc.GetPathDescription().GetTablePartitions();
-            UNIT_ASSERT(!tablePartitions.empty());
-            writtenToMainTable = mainTable.ExpectedRecords == CountRows(runtime, mainTable, tablePartitions);
-        }
+        auto mainTableRows = CountRows(runtime, mainTable.Path);
+        bool writtenToMainTable = (mainTable.ExpectedRecords == mainTableRows);
 
         if (writtenToMainTable) {
-            auto tableDesc = DescribePrivatePath(runtime, indexTable.Path, true, true);
-            const auto& tablePartitions = tableDesc.GetPathDescription().GetTablePartitions();
-            UNIT_ASSERT(!tablePartitions.empty());
-
             int i = 0;
             while (++i < 10) {
                 runtime.SimulateSleep(TDuration::Seconds(1));
-                if (indexTable.ExpectedRecords == CountRows(runtime, indexTable, tablePartitions)) {
+                if (indexTable.ExpectedRecords == CountRows(runtime, indexTable.Path)) {
                     break;
                 }
             }
diff --git a/ydb/core/tx/schemeshard/ut_index/ya.make b/ydb/core/tx/schemeshard/ut_index/ya.make
index 6bdda0b6204..ddd365aeee7 100644
--- a/ydb/core/tx/schemeshard/ut_index/ya.make
+++ b/ydb/core/tx/schemeshard/ut_index/ya.make
@@ -16,7 +16,6 @@ PEERDIR(
     ydb/core/scheme
     ydb/core/testlib/default
     ydb/core/tx/schemeshard/ut_helpers
-    ydb/public/lib/deprecated/kicli
 )
 
 SRCS(
diff --git a/ydb/core/tx/schemeshard/ut_index_build/ut_vector_index_build.cpp b/ydb/core/tx/schemeshard/ut_index_build/ut_vector_index_build.cpp
index fe5c69c88b0..f2d8530bd57 100644
--- a/ydb/core/tx/schemeshard/ut_index_build/ut_vector_index_build.cpp
+++ b/ydb/core/tx/schemeshard/ut_index_build/ut_vector_index_build.cpp
@@ -8,7 +8,6 @@
 #include <ydb/core/tx/datashard/datashard.h>
 #include <ydb/core/metering/metering.h>
 
-#include <ydb/public/lib/deprecated/kicli/kicli.h>
 #include <ydb/public/sdk/cpp/include/ydb-cpp-sdk/client/table/table.h>
 
 using namespace NKikimr;
@@ -177,15 +176,7 @@ Y_UNIT_TEST_SUITE (VectorIndexBuildTest) {
 
         // Check row count in the posting table
         {
-            auto indexDesc = DescribePath(runtime, tenantSchemeShard, "/MyRoot/ServerLessDB/Table/index1/indexImplPostingTable", true, true, true);
-            auto parts = indexDesc.GetPathDescription().GetTablePartitions();
-            ui32 rows = 0;
-            for (const auto & x: parts) {
-                auto result = ReadTable(runtime, x.GetDatashardId(), "indexImplPostingTable",
-                    {NKikimr::NTableIndex::NTableVectorKmeansTreeIndex::ParentColumn, "key"}, {"key"});
-                auto value = NClient::TValue::Create(result);
-                rows += value["Result"]["List"].Size();
-            }
+            auto rows = CountRows(runtime, tenantSchemeShard, "/MyRoot/ServerLessDB/Table/index1/indexImplPostingTable");
             Cerr << "... posting table contains " << rows << " rows" << Endl;
             UNIT_ASSERT_VALUES_EQUAL(rows, 200);
         }