Fix TKeyValueTracingTest unit tests (#26430)

Author: serbel324

.github/config/muted_ya.txt

@@ -18,10 +18,6 @@ ydb/core/fq/libs/row_dispatcher/format_handler/ut TestFormatHandler.WatermarkWhe
 ydb/core/fq/libs/row_dispatcher/format_handler/ut TestFormatHandler.WatermarkWhereFalse
 ydb/core/kafka_proxy/ut KafkaProtocol.TopicsCompactionSwitchOnAndOff
 ydb/core/kafka_proxy/ut unittest.sole chunk
-ydb/core/keyvalue/ut_trace TKeyValueTracingTest.ReadHuge
-ydb/core/keyvalue/ut_trace TKeyValueTracingTest.ReadSmall
-ydb/core/keyvalue/ut_trace TKeyValueTracingTest.WriteHuge
-ydb/core/keyvalue/ut_trace TKeyValueTracingTest.WriteSmall
 ydb/core/kqp/ut/close_with_load KqpQueryService.CloseSessionsWithLoad
 ydb/core/kqp/ut/close_with_load KqpService.CloseSessionsWithLoad
 ydb/core/kqp/ut/close_with_load unittest.[*/*] chunk

ydb/core/keyvalue/keyvalue_ut_trace.cpp

@@ -100,14 +100,18 @@ void TestOneWrite(TString value, TVector<TString> &&expectedTraceVariants) {
     env.DoKVRequest(CreateWrite("key", std::move(value)));
 
     UNIT_ASSERT(env.WilsonUploader->BuildTraceTrees());
-    UNIT_ASSERT_VALUES_EQUAL(env.WilsonUploader->Traces.size(), 1);
-    auto& trace = env.WilsonUploader->Traces.begin()->second;
+    // we expect one trace from DSProxy and one from KV-tablet
+    UNIT_ASSERT_VALUES_EQUAL(env.WilsonUploader->Traces.size(), 2);
+
+    bool oneIsGood = std::any_of(env.WilsonUploader->Traces.begin(), env.WilsonUploader->Traces.end(),
+        [&](const auto& v) {
+            const auto& [_, trace] = v;
+            return std::any_of(expectedTraceVariants.begin(), expectedTraceVariants.end(),
+                    [&](const TString& var) { return var == trace.ToString(); });
+        }
+    );
 
-    bool found = false;
-    for (const TString &expectedTrace : expectedTraceVariants) {
-        found |= trace.ToString() == expectedTrace;
-    }
-    UNIT_ASSERT_C(found, trace.ToString());
+    UNIT_ASSERT(oneIsGood);
 }
 
 void TestOneRead(TString value, TString expectedTrace) {
@@ -121,10 +125,16 @@ void TestOneRead(TString value, TString expectedTrace) {
     UNIT_ASSERT_EQUAL(response.value(), value);
 
     UNIT_ASSERT(env.WilsonUploader->BuildTraceTrees());
-    UNIT_ASSERT_EQUAL(env.WilsonUploader->Traces.size(), 1);
-    auto& trace = env.WilsonUploader->Traces.begin()->second;
+    // we expect one trace from DSProxy and one from KV-tablet
+    UNIT_ASSERT_VALUES_EQUAL(env.WilsonUploader->Traces.size(), 2);
 
-    UNIT_ASSERT_EQUAL(trace.ToString(), expectedTrace);
+    bool oneIsGood = std::any_of(env.WilsonUploader->Traces.begin(), env.WilsonUploader->Traces.end(),
+        [&](const auto& v) {
+            const auto& [_, trace] = v;
+            return trace.ToString() == expectedTrace;
+        }
+    );
+    UNIT_ASSERT(oneIsGood);
 }
 
 Y_UNIT_TEST_SUITE(TKeyValueTracingTest) {
@@ -146,6 +156,7 @@ TVector<std::string> PDiskLogWriteVariants = {
     "(PDisk.LogWrite -> "
         "[(PDisk.InScheduler) , (PDisk.InBlockDevice)]"
     ")",
+    "(PDisk.LogWrite)",
 };
 
 TVector<std::string> TabletWriteLogTempaltes = {
@@ -224,14 +235,10 @@ Y_UNIT_TEST(WriteSmall) {
     TVector<TString> canons = MakeCanons(
         "(KeyValue.Intermediate -> "
             "[(KeyValue.StorageRequest -> "
-                "[(DSProxy.Put -> "
-                    "[(Backpressure.InFlight -> "
-                        "[(VDisk.Log.Put)]"
-                    ")]"
-                ")]"
+                "[(DSProxy.Put)]"
             ") , "
             "(Tablet.Transaction -> "
-                "[(Tablet.Transaction.Execute) , {TABLET_LOG_WRITE}]"
+                "[(Tablet.Transaction.Execute) , {TABLET_LOG_WRITE} , (Tablet.Transaction.Complete)]"
             ")]"
         ")"
     );
@@ -242,34 +249,23 @@ Y_UNIT_TEST(WriteHuge) {
     TVector<TString> canons = MakeCanons(
         "(KeyValue.Intermediate -> "
             "[(KeyValue.StorageRequest -> "
-                "[(DSProxy.Put -> "
-                    "[(Backpressure.InFlight -> "
-                        "[(VDisk.HullHugeBlobChunkAllocator) , "
-                        "(VDisk.HullHugeKeeper.InWaitQueue -> "
-                            "[(VDisk.HugeBlobKeeper.Write -> "
-                                "[(VDisk.Log.PutHuge)]"
-                            ")]"
-                        ")]"
-                    ")]"
-                ")]"
+                "[(DSProxy.Put)]"
             ") , "
             "(Tablet.Transaction -> "
-                "[(Tablet.Transaction.Execute) , {TABLET_LOG_WRITE}]"
+                "[(Tablet.Transaction.Execute) , {TABLET_LOG_WRITE} , (Tablet.Transaction.Complete)]"
             ")]"
         ")"
     );
     TestOneWrite(HugeValue, std::move(canons));
 }
 
 Y_UNIT_TEST(ReadSmall) {
-    TString canon = "(KeyValue.Intermediate -> [(KeyValue.StorageReadRequest -> [(DSProxy.Get -> [(Backpressure.InFlight -> "
-        "[(VDisk.LevelIndexExtremeQueryViaBatcherMergeData)])])])])";
+    TString canon = "(KeyValue.Intermediate -> [(KeyValue.StorageReadRequest -> [(DSProxy.Get)])])";
     TestOneRead(SmallValue, std::move(canon));
 }
 
 Y_UNIT_TEST(ReadHuge) {
-    TString canon = "(KeyValue.Intermediate -> [(KeyValue.StorageReadRequest -> [(DSProxy.Get -> [(Backpressure.InFlight -> "
-        "[(VDisk.LevelIndexExtremeQueryViaBatcherMergeData -> [(VDisk.Query.ReadBatcher)])])])])])";
+    TString canon = "(KeyValue.Intermediate -> [(KeyValue.StorageReadRequest -> [(DSProxy.Get)])])";
     TestOneRead(HugeValue, std::move(canon));
 }