Improve read table test coverage: ut to read one row from table by StreamReadTable call (#12118)

1 files changed, 41 insertions, 14 deletions

diff --git a/ydb/services/ydb/ydb_table_ut.cpp b/ydb/services/ydb/ydb_table_ut.cpp
index b4d004e99a..6cdad26902 100644
--- a/ydb/services/ydb/ydb_table_ut.cpp
+++ b/ydb/services/ydb/ydb_table_ut.cpp
@@ -1787,7 +1787,13 @@ R"___(<main>: Error: Transaction not found: , code: 2015
         UseSnapshot,
     };
 
-    void TestReadTableMultiShard(EReadTableMultiShardMode mode, bool wholeTable) {
+    enum class EReadTableRangeMode {
+        OneRow,
+        TwoRows,
+        WholeTable
+    };
+
+    void TestReadTableMultiShard(EReadTableMultiShardMode mode, EReadTableRangeMode rangeMode) {
         TKikimrWithGrpcAndRootSchema server;
         ui16 grpc = server.GetPort();
 
@@ -1833,15 +1839,23 @@ R"___(<main>: Error: Transaction not found: , code: 2015
             .EndTuple();
 
         TValueBuilder valueTo;
-        valueTo.BeginTuple()
-            .AddElement()
-                .OptionalUint32(1000000000u)
-            .AddElement()
-                .OptionalUint64(2000000000u)
-            .EndTuple();
+
+        if (rangeMode == EReadTableRangeMode::OneRow) {
+            valueTo.BeginTuple()
+                .AddElement()
+                    .OptionalUint32(1u)
+                .EndTuple();
+        } else {
+            valueTo.BeginTuple()
+                .AddElement()
+                    .OptionalUint32(1000000000u)
+                .AddElement()
+                    .OptionalUint64(2000000000u)
+                .EndTuple();
+        }
 
         TReadTableSettings readTableSettings =
-            wholeTable ? TReadTableSettings().Ordered() :
+            rangeMode == EReadTableRangeMode::WholeTable ? TReadTableSettings().Ordered() :
             TReadTableSettings()
                 .Ordered()
                 .From(TKeyBound::Inclusive(valueFrom.Build()))
@@ -1865,7 +1879,7 @@ R"___(<main>: Error: Transaction not found: , code: 2015
         TVector<TRows> expected;
         expected.push_back({1u, 1u, "One"});
         expected.push_back({1000000000u, 2u, "Two"});
-        if (wholeTable) {
+        if (rangeMode == EReadTableRangeMode::WholeTable) {
             expected.push_back({4294967295u, 4u, "Last"});
         }
         int row = 0;
@@ -1895,26 +1909,39 @@ R"___(<main>: Error: Transaction not found: , code: 2015
                 UNIT_ASSERT_VALUES_EQUAL(val, exp.Value);
             }
         }
-        UNIT_ASSERT_VALUES_EQUAL(row, wholeTable ? 3 : 2);
+        switch (rangeMode) {
+            case EReadTableRangeMode::OneRow:
+                UNIT_ASSERT_VALUES_EQUAL(row, 1);
+                break;
+            case EReadTableRangeMode::TwoRows:
+                UNIT_ASSERT_VALUES_EQUAL(row, 2);
+                break;
+            case EReadTableRangeMode::WholeTable:
+                UNIT_ASSERT_VALUES_EQUAL(row, 3);
+        }
 
         // Attempt to call ReadNext on finished iterator causes ContractViolation
         UNIT_ASSERT_EXCEPTION(it.ReadNext().GetValueSync().EOS(), NYdb::TContractViolation);
     }
 
     Y_UNIT_TEST(TestReadTableMultiShard) {
-        TestReadTableMultiShard(EReadTableMultiShardMode::Normal, false);
+        TestReadTableMultiShard(EReadTableMultiShardMode::Normal, EReadTableRangeMode::TwoRows);
     }
 
     Y_UNIT_TEST(TestReadTableMultiShardUseSnapshot) {
-        TestReadTableMultiShard(EReadTableMultiShardMode::UseSnapshot, false);
+        TestReadTableMultiShard(EReadTableMultiShardMode::UseSnapshot, EReadTableRangeMode::TwoRows);
     }
 
     Y_UNIT_TEST(TestReadTableMultiShardWholeTable) {
-        TestReadTableMultiShard(EReadTableMultiShardMode::Normal, true);
+        TestReadTableMultiShard(EReadTableMultiShardMode::Normal, EReadTableRangeMode::WholeTable);
     }
 
     Y_UNIT_TEST(TestReadTableMultiShardWholeTableUseSnapshot) {
-        TestReadTableMultiShard(EReadTableMultiShardMode::UseSnapshot, true);
+        TestReadTableMultiShard(EReadTableMultiShardMode::UseSnapshot, EReadTableRangeMode::WholeTable);
+    }
+
+    Y_UNIT_TEST(TestReadTableMultiShardOneRow) {
+        TestReadTableMultiShard(EReadTableMultiShardMode::Normal, EReadTableRangeMode::OneRow);
     }
 
     void TestReadTableMultiShardWithDescribe(bool rowLimit) {