Refactor topic tests

Author: Gazizonoki

src/client/topic/ut/topic_to_table_ut.cpp

@@ -3,21 +3,25 @@
 #include <ydb-cpp-sdk/client/topic/client.h>
 #include <ydb-cpp-sdk/client/table/table.h>
 #include <src/client/persqueue_public/ut/ut_utils/ut_utils.h>
+
 #include <ydb/core/cms/console/console.h>
 #include <ydb/core/keyvalue/keyvalue_events.h>
 #include <ydb/core/persqueue/key.h>
 #include <ydb/core/persqueue/blob.h>
-
+#include <ydb/core/persqueue/events/global.h>
 #include <ydb/core/tx/long_tx_service/public/events.h>
 
-#include <library/cpp/logger/stream.h>
+#include <ydb/core/persqueue/ut/common/autoscaling_ut_common.h>
 
+#include <library/cpp/logger/stream.h>
 #include <library/cpp/testing/unittest/registar.h>
 
 namespace NYdb::NTopic::NTests {
 
 const auto TEST_MESSAGE_GROUP_ID_1 = TEST_MESSAGE_GROUP_ID + "_1";
 const auto TEST_MESSAGE_GROUP_ID_2 = TEST_MESSAGE_GROUP_ID + "_2";
+const auto TEST_MESSAGE_GROUP_ID_3 = TEST_MESSAGE_GROUP_ID + "_3";
+const auto TEST_MESSAGE_GROUP_ID_4 = TEST_MESSAGE_GROUP_ID + "_4";
 
 Y_UNIT_TEST_SUITE(TxUsage) {
 
@@ -79,9 +83,16 @@ class TFixture : public NUnitTest::TBaseFixture {
                      const TString& consumer = TEST_CONSUMER,
                      size_t partitionCount = 1,
                      std::optional<size_t> maxPartitionCount = std::nullopt);
-    void DescribeTopic(const TString& path);
+    TTopicDescription DescribeTopic(const TString& path);
 
-    void AddConsumer(const TString& topic, const TVector<TString>& consumers);
+    void AddConsumer(const TString& topicPath, const TVector<TString>& consumers);
+    void AlterAutoPartitioning(const TString& topicPath,
+                               ui64 minActivePartitions,
+                               ui64 maxActivePartitions,
+                               EAutoPartitioningStrategy strategy,
+                               TDuration stabilizationWindow,
+                               ui64 downUtilizationPercent,
+                               ui64 upUtilizationPercent);
 
     void WriteToTopicWithInvalidTxId(bool invalidTxId);
 
@@ -185,6 +196,24 @@ class TFixture : public NUnitTest::TBaseFixture {
                                         const TString& consumerName,
                                         size_t count);
 
+    struct TAvgWriteBytes {
+        ui64 PerSec = 0;
+        ui64 PerMin = 0;
+        ui64 PerHour = 0;
+        ui64 PerDay = 0;
+    };
+
+    TAvgWriteBytes GetAvgWriteBytes(const TString& topicPath,
+                                    ui32 partitionId);
+
+    void CheckAvgWriteBytes(const TString& topicPath,
+                            ui32 partitionId,
+                            size_t minSize, size_t maxSize);
+
+    void SplitPartition(const TString& topicPath,
+                        ui32 partitionId,
+                        const TString& boundary);
+
 private:
     template<class E>
     E ReadEvent(TTopicReadSessionPtr reader, NTable::TTransaction& tx);
@@ -213,6 +242,8 @@ class TFixture : public NUnitTest::TBaseFixture {
 
     THashMap<std::pair<TString, TString>, TTopicWriteSessionContext> TopicWriteSessions;
     THashMap<TString, TTopicReadSessionPtr> TopicReadSessions;
+
+    ui64 SchemaTxId = 1000;
 };
 
 TFixture::TTableRecord::TTableRecord(const TString& key, const TString& value) :
@@ -225,6 +256,7 @@ void TFixture::SetUp(NUnitTest::TTestContext&)
 {
     NKikimr::Tests::TServerSettings settings = TTopicSdkTestSetup::MakeServerSettings();
     settings.SetEnableTopicServiceTx(true);
+    settings.SetEnableTopicSplitMerge(true);
     settings.SetEnablePQConfigTransactionsAtSchemeShard(true);
 
     Setup = std::make_unique<TTopicSdkTestSetup>(TEST_CASE_NAME, settings);
@@ -405,7 +437,7 @@ void TFixture::CreateTopic(const TString& path,
     Setup->CreateTopic(path, consumer, partitionCount, maxPartitionCount);
 }
 
-void TFixture::AddConsumer(const TString& path,
+void TFixture::AddConsumer(const TString& topicPath,
                            const TVector<TString>& consumers)
 {
     NTopic::TTopicClient client(GetDriver());
@@ -415,13 +447,41 @@ void TFixture::AddConsumer(const TString& path,
         settings.BeginAddConsumer(consumer);
     }
 
-    auto result = client.AlterTopic(path, settings).GetValueSync();
+    auto result = client.AlterTopic(topicPath, settings).GetValueSync();
     UNIT_ASSERT_C(result.IsSuccess(), result.GetIssues().ToString());
 }
 
-void TFixture::DescribeTopic(const TString& path)
+void TFixture::AlterAutoPartitioning(const TString& topicPath,
+                                     ui64 minActivePartitions,
+                                     ui64 maxActivePartitions,
+                                     EAutoPartitioningStrategy strategy,
+                                     TDuration stabilizationWindow,
+                                     ui64 downUtilizationPercent,
+                                     ui64 upUtilizationPercent)
 {
-    Setup->DescribeTopic(path);
+    NTopic::TTopicClient client(GetDriver());
+    NTopic::TAlterTopicSettings settings;
+
+    settings
+        .BeginAlterPartitioningSettings()
+            .MinActivePartitions(minActivePartitions)
+            .MaxActivePartitions(maxActivePartitions)
+            .BeginAlterAutoPartitioningSettings()
+                .Strategy(strategy)
+                .StabilizationWindow(stabilizationWindow)
+                .DownUtilizationPercent(downUtilizationPercent)
+                .UpUtilizationPercent(upUtilizationPercent)
+            .EndAlterAutoPartitioningSettings()
+        .EndAlterTopicPartitioningSettings()
+        ;
+
+    auto result = client.AlterTopic(topicPath, settings).GetValueSync();
+    UNIT_ASSERT_C(result.IsSuccess(), result.GetIssues().ToString());
+}
+
+TTopicDescription TFixture::DescribeTopic(const TString& path)
+{
+    return Setup->DescribeTopic(path);
 }
 
 const TDriver& TFixture::GetDriver() const
@@ -545,11 +605,11 @@ Y_UNIT_TEST_F(Offsets_Cannot_Be_Promoted_When_Reading_In_A_Transaction, TFixture
     UNIT_ASSERT_EXCEPTION(ReadMessage(reader, {.Tx = tx, .CommitOffsets = true}), yexception);
 }
 
-//Y_UNIT_TEST_F(WriteToTopic_Invalid_Session, TFixture)
-//{
-//    WriteToTopicWithInvalidTxId(false);
-//}
-//
+Y_UNIT_TEST_F(WriteToTopic_Invalid_Session, TFixture)
+{
+    WriteToTopicWithInvalidTxId(false);
+}
+
 //Y_UNIT_TEST_F(WriteToTopic_Invalid_Tx, TFixture)
 //{
 //    WriteToTopicWithInvalidTxId(true);
@@ -1021,6 +1081,34 @@ void TFixture::Read_Exactly_N_Messages_From_Topic(const TString& topicPath,
     UNIT_ASSERT_VALUES_EQUAL(count, limit);
 }
 
+auto TFixture::GetAvgWriteBytes(const TString& topicName,
+                                ui32 partitionId) -> TAvgWriteBytes
+{
+    auto& runtime = Setup->GetRuntime();
+    TActorId edge = runtime.AllocateEdgeActor();
+    ui64 tabletId = GetTopicTabletId(edge, "/Root/" + topicName, partitionId);
+
+    runtime.SendToPipe(tabletId, edge, new NKikimr::TEvPersQueue::TEvStatus());
+    auto response = runtime.GrabEdgeEvent<NKikimr::TEvPersQueue::TEvStatusResponse>();
+
+    UNIT_ASSERT_VALUES_EQUAL(tabletId, response->Record.GetTabletId());
+
+    TAvgWriteBytes result;
+
+    for (size_t i = 0; i < response->Record.PartResultSize(); ++i) {
+        const auto& partition = response->Record.GetPartResult(i);
+        if (partition.GetPartition() == static_cast<int>(partitionId)) {
+            result.PerSec = partition.GetAvgWriteSpeedPerSec();
+            result.PerMin = partition.GetAvgWriteSpeedPerMin();
+            result.PerHour = partition.GetAvgWriteSpeedPerHour();
+            result.PerDay = partition.GetAvgWriteSpeedPerDay();
+            break;
+        }
+    }
+
+    return result;
+}
+
 Y_UNIT_TEST_F(WriteToTopic_Demo_1, TFixture)
 {
     CreateTopic("topic_A");
@@ -2265,6 +2353,150 @@ Y_UNIT_TEST_F(WriteToTopic_Demo_44, TFixture)
     Read_Exactly_N_Messages_From_Topic("topic_A", TEST_CONSUMER, 100);
 }
 
+void TFixture::CheckAvgWriteBytes(const TString& topicPath,
+                                  ui32 partitionId,
+                                  size_t minSize, size_t maxSize)
+{
+#define UNIT_ASSERT_AVGWRITEBYTES(v, minSize, maxSize) \
+    UNIT_ASSERT_LE_C(minSize, v, ", actual " << minSize << " > " << v); \
+    UNIT_ASSERT_LE_C(v, maxSize, ", actual " << v << " > " << maxSize);
+
+    auto avgWriteBytes = GetAvgWriteBytes(topicPath, partitionId);
+
+    UNIT_ASSERT_AVGWRITEBYTES(avgWriteBytes.PerSec, minSize, maxSize);
+    UNIT_ASSERT_AVGWRITEBYTES(avgWriteBytes.PerMin, minSize, maxSize);
+    UNIT_ASSERT_AVGWRITEBYTES(avgWriteBytes.PerHour, minSize, maxSize);
+    UNIT_ASSERT_AVGWRITEBYTES(avgWriteBytes.PerDay, minSize, maxSize);
+
+#undef UNIT_ASSERT_AVGWRITEBYTES
+}
+
+void TFixture::SplitPartition(const TString& topicName,
+                              ui32 partitionId,
+                              const TString& boundary)
+{
+    NKikimr::NPQ::NTest::SplitPartition(Setup->GetRuntime(),
+                                        ++SchemaTxId,
+                                        topicName,
+                                        partitionId,
+                                        boundary);
+}
+
+Y_UNIT_TEST_F(WriteToTopic_Demo_45, TFixture)
+{
+    // Writing to a topic in a transaction affects the `AvgWriteBytes` indicator
+    CreateTopic("topic_A", TEST_CONSUMER, 2);
+
+    auto session = CreateTableSession();
+    auto tx = BeginTx(session);
+
+    TString message(1'000, 'x');
+
+    WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID_1, message, &tx, 0);
+    WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID_1, message, &tx, 0);
+
+    WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID_2, message, &tx, 1);
+
+    CommitTx(tx, EStatus::SUCCESS);
+
+    size_t minSize = (message.size() + TEST_MESSAGE_GROUP_ID_1.size()) * 2;
+    size_t maxSize = minSize + 200;
+
+    CheckAvgWriteBytes("topic_A", 0, minSize, maxSize);
+
+    minSize = (message.size() + TEST_MESSAGE_GROUP_ID_2.size());
+    maxSize = minSize + 200;
+
+    CheckAvgWriteBytes("topic_A", 1, minSize, maxSize);
+}
+
+Y_UNIT_TEST_F(WriteToTopic_Demo_46, TFixture)
+{
+    // The `split` operation of the topic partition affects the writing in the transaction.
+    // The transaction commit should fail with an error
+    CreateTopic("topic_A", TEST_CONSUMER, 2, 10);
+
+    auto session = CreateTableSession();
+    auto tx = BeginTx(session);
+
+    TString message(1'000, 'x');
+
+    WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID_1, message, &tx, 0);
+    WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID_1, message, &tx, 0);
+
+    WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID_2, message, &tx, 1);
+
+    WaitForAcks("topic_A", TEST_MESSAGE_GROUP_ID_2);
+
+    SplitPartition("topic_A", 1, "\xC0");
+
+    WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID_2, message, &tx, 1);
+
+    CommitTx(tx, EStatus::ABORTED);
+}
+
+Y_UNIT_TEST_F(WriteToTopic_Demo_47, TFixture)
+{
+    // The `split` operation of the topic partition does not affect the reading in the transaction.
+    CreateTopic("topic_A", TEST_CONSUMER, 2, 10);
+
+    TString message(1'000, 'x');
+
+    WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID_1, message, nullptr, 0);
+    WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID_1, message, nullptr, 0);
+
+    WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID_2, message, nullptr, 1);
+    WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID_2, message, nullptr, 1);
+
+    WaitForAcks("topic_A", TEST_MESSAGE_GROUP_ID_1);
+    WaitForAcks("topic_A", TEST_MESSAGE_GROUP_ID_2);
+
+    SplitPartition("topic_A", 1, "\xC0");
+
+    auto session = CreateTableSession();
+    auto tx = BeginTx(session);
+
+    auto messages = ReadFromTopic("topic_A", TEST_CONSUMER, TDuration::Seconds(2), &tx, 0);
+    UNIT_ASSERT_VALUES_EQUAL(messages.size(), 2);
+
+    CloseTopicReadSession("topic_A", TEST_CONSUMER);
+
+    messages = ReadFromTopic("topic_A", TEST_CONSUMER, TDuration::Seconds(2), &tx, 1);
+    UNIT_ASSERT_VALUES_EQUAL(messages.size(), 2);
+
+    CommitTx(tx, EStatus::SUCCESS);
+}
+
+Y_UNIT_TEST_F(WriteToTopic_Demo_48, TFixture)
+{
+    // the commit of a transaction affects the split of the partition
+    CreateTopic("topic_A", TEST_CONSUMER, 2, 10);
+    AlterAutoPartitioning("topic_A", 2, 10, EAutoPartitioningStrategy::ScaleUp, TDuration::Seconds(2), 1, 2);
+
+    auto session = CreateTableSession();
+    auto tx = BeginTx(session);
+
+    TString message(1_MB, 'x');
+
+    WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID_1, message, &tx, 0);
+    WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID_1, message, &tx, 0);
+    WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID_3, message, &tx, 0);
+    WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID_3, message, &tx, 0);
+
+    WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID_2, message, &tx, 1);
+    WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID_2, message, &tx, 1);
+    WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID_4, message, &tx, 1);
+    WriteToTopic("topic_A", TEST_MESSAGE_GROUP_ID_4, message, &tx, 1);
+
+    CommitTx(tx, EStatus::SUCCESS);
+
+    Sleep(TDuration::Seconds(5));
+
+    auto topicDescription = DescribeTopic("topic_A");
+
+    UNIT_ASSERT_GT(topicDescription.GetTotalPartitionsCount(), 2);
+}
+
 }
 
 }

src/client/topic/ut/ut_utils/topic_sdk_test_setup.cpp

@@ -52,7 +52,7 @@ void TTopicSdkTestSetup::CreateTopic(const TString& path, const TString& consume
     Server.WaitInit(path);
 }
 
-void TTopicSdkTestSetup::DescribeTopic(const TString& path)
+TTopicDescription TTopicSdkTestSetup::DescribeTopic(const TString& path)
 {
     TTopicClient client(MakeDriver());
 
@@ -62,6 +62,8 @@ void TTopicSdkTestSetup::DescribeTopic(const TString& path)
 
     auto status = client.DescribeTopic(path, settings).GetValueSync();
     UNIT_ASSERT(status.IsSuccess());
+
+    return status.GetTopicDescription();
 }
 
 TString TTopicSdkTestSetup::GetEndpoint() const {

src/client/topic/ut/ut_utils/topic_sdk_test_setup.h

@@ -21,7 +21,7 @@ class TTopicSdkTestSetup {
     void CreateTopicWithAutoscale(const TString& path = TString{TEST_TOPIC}, const TString& consumer = TEST_CONSUMER, size_t partitionCount = 1,
                      size_t maxPartitionCount = 100);
 
-    void DescribeTopic(const TString& path = TString{TEST_TOPIC});
+    TTopicDescription DescribeTopic(const TString& path = TString{TEST_TOPIC});
 
     TString GetEndpoint() const;
     TString GetTopicPath(const TString& name = TString{TEST_TOPIC}) const;