speed up test

2 files changed, 40 insertions, 22 deletions

diff --git a/ydb/core/persqueue/ut/common/pq_ut_common.h b/ydb/core/persqueue/ut/common/pq_ut_common.h
index 2d187d87f8..8dc97c7812 100644
--- a/ydb/core/persqueue/ut/common/pq_ut_common.h
+++ b/ydb/core/persqueue/ut/common/pq_ut_common.h
@@ -25,15 +25,46 @@ constexpr ui32 NUM_WRITES = PlainOrSoSlow(50, 1);
 
 void FillPQConfig(NKikimrPQ::TPQConfig& pqConfig, const TString& dbRoot, bool isFirstClass);
 
+enum EventKing {
+    TabletPipe,
+    NPDisk,
+    KeyValue,
+    PQ
+};
+
 class TInitialEventsFilter : TNonCopyable {
-    bool IsDone;
+    std::unordered_set<TString> Events;
 public:
-    TInitialEventsFilter()
-        : IsDone(false)
-    {}
+    TInitialEventsFilter() = default;
+
+    TTestActorRuntime::TEventFilter Prepare(const std::unordered_set<EventKing>& eventKings = {TabletPipe, NPDisk, KeyValue, PQ}, 
+                         const std::unordered_set<TString>& eventTypeNames = {}) {
+        Events.clear();
+
+        if (eventKings.contains(TabletPipe)) {
+            Events.insert("NKikimr::TEvTabletPipe::TEvClientConnected");
+            Events.insert("NKikimr::TEvTabletPipe::TEvClientDestroyed");
+            Events.insert("NKikimr::TEvTabletPipe::TEvServerConnected");
+        }
+        if (eventKings.contains(NPDisk)) {
+            Events.insert("NKikimr::NPDisk::TEvLog");
+            Events.insert("NKikimr::NPDisk::TEvLogResult");
+        }
+        if (eventKings.contains(KeyValue)) {
+            Events.insert("NKikimr::TEvKeyValue::TEvCollect");
+            Events.insert("NKikimr::TEvKeyValue::TEvCompleteGC");
+            Events.insert("NKikimr::TEvKeyValue::TEvIntermediate");
+            Events.insert("NKikimr::TEvKeyValue::TEvPartialCompleteGC");
+        }
+        if (eventKings.contains(PQ)) {
+            Events.insert("NKikimr::TEvPQ::TEvPartitionLabeledCounters");
+            Events.insert("NKikimr::TEvPQ::TEvProxyResponse");
+        }
+
+        for(const auto& v : eventTypeNames) {
+            Events.insert(v);
+        }
 
-    TTestActorRuntime::TEventFilter Prepare() {
-        IsDone = false;
         return [&](TTestActorRuntimeBase& runtime, TAutoPtr<IEventHandle>& event) {
             return (*this)(runtime, event);
         };
@@ -41,12 +72,12 @@ public:
 
     bool operator()(TTestActorRuntimeBase& runtime, TAutoPtr<IEventHandle>& event) {
         Y_UNUSED(runtime);
-        Y_UNUSED(event);
 
-        return false;
+        return Events.contains(event->GetTypeName());
     }
 };
 
+
 struct TTestContext {
     const TTabletTypes::EType PQTabletType = TTabletTypes::PersQueue;
     const TTabletTypes::EType BalancerTabletType = TTabletTypes::PersQueueReadBalancer;
diff --git a/ydb/core/persqueue/ut/pq_ut.cpp b/ydb/core/persqueue/ut/pq_ut.cpp
index 1149f285d3..ec85e940c6 100644
--- a/ydb/core/persqueue/ut/pq_ut.cpp
+++ b/ydb/core/persqueue/ut/pq_ut.cpp
@@ -1288,20 +1288,7 @@ Y_UNIT_TEST(TestWriteSplit) {
 Y_UNIT_TEST(TestLowWatermark) {
     TTestContext tc;
     RunTestWithReboots(tc.TabletIds, [&]() {
-        return [](TTestActorRuntimeBase&, TAutoPtr<IEventHandle>& event) -> bool {
-            const auto name = event->GetTypeName();
-            return name == "NKikimr::TEvTabletPipe::TEvClientConnected" ||
-                   name == "NKikimr::TEvTabletPipe::TEvClientDestroyed" ||
-                   name == "NKikimr::TEvTabletPipe::TEvServerConnected" ||
-                   name == "NKikimr::NPDisk::TEvLog" ||
-                   name == "NKikimr::NPDisk::TEvLogResult" ||
-                   name == "NKikimr::TEvKeyValue::TEvCollect" ||
-                   name == "NKikimr::TEvKeyValue::TEvCompleteGC" ||
-                   name == "NKikimr::TEvKeyValue::TEvIntermediate" ||
-                   name == "NKikimr::TEvKeyValue::TEvPartialCompleteGC" ||
-                   name == "NKikimr::TEvPQ::TEvPartitionLabeledCounters" ||
-                   name == "NKikimr::TEvPQ::TEvProxyResponse";
-        };
+        return tc.InitialEventsFilter.Prepare();
     }, [&](const TString& dispatchName, std::function<void(TTestActorRuntime&)> setup, bool& activeZone) {
         TFinalizer finalizer(tc);
         tc.Prepare(dispatchName, setup, activeZone);