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

Author: vitalif

ydb/core/tx/schemeshard/ut_helpers/helpers.cpp

@@ -2628,6 +2628,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 };

ydb/core/tx/schemeshard/ut_helpers/helpers.h

@@ -653,6 +653,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);
 

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
 )

ydb/core/tx/schemeshard/ut_index/ut_async_index.cpp

@@ -126,37 +126,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;
                 }
             }

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(

ydb/core/tx/schemeshard/ut_index_build/ut_vector_index_build.cpp

@@ -177,15 +177,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);
         }