Make it easier to use IEventHandle as an unfiltered event type in tes…

…ts (#8723)
ydb-platform · Sep 4, 2024 · 93bd594 · 93bd594
1 parent f8bebff
commit 93bd594
Show file tree

Hide file tree

Showing 2 changed files with 49 additions and 18 deletions.
diff --git a/ydb/library/actors/core/event.h b/ydb/library/actors/core/event.h
@@ -67,6 +67,9 @@ namespace NActors {
             }
         };
 
+    public:
+        typedef TAutoPtr<IEventHandle> TPtr;
+
     public:
         template <typename TEv>
         inline TEv* CastAsLocal() const noexcept {
@@ -349,6 +352,10 @@ namespace NActors {
     template <typename TEventType>
     class TEventHandle: public IEventHandle {
         TEventHandle(); // we never made instance of TEventHandle
+
+    public:
+        typedef TAutoPtr<TEventHandle<TEventType>> TPtr;
+
     public:
         TEventType* Get() {
             return IEventHandle::Get<TEventType>();
@@ -371,6 +378,6 @@ namespace NActors {
         // still abstract
 
         typedef TEventHandle<TEventType> THandle;
-        typedef TAutoPtr<THandle> TPtr;
+        typedef typename THandle::TPtr TPtr;
     };
 }
diff --git a/ydb/library/actors/testlib/test_runtime.h b/ydb/library/actors/testlib/test_runtime.h
@@ -203,6 +203,36 @@ namespace NActors {
         }
     };
 
+    /**
+     * Allows customizing behavior based on the event type
+     */
+    template<class TEvType>
+    struct TTestEventObserverTraits {
+        static bool Match(IEventHandle::TPtr& ev) noexcept {
+            return ev->GetTypeRewrite() == TEvType::EventType;
+        }
+
+        static typename TEvType::TPtr& Convert(IEventHandle::TPtr& ev) noexcept {
+            return reinterpret_cast<typename TEvType::TPtr&>(ev);
+        }
+    };
+
+    template<>
+    struct TTestEventObserverTraits<IEventHandle> {
+        static constexpr bool Match(IEventHandle::TPtr&) noexcept {
+            return true;
+        }
+
+        static constexpr IEventHandle::TPtr& Convert(IEventHandle::TPtr& ev) noexcept {
+            return ev;
+        }
+    };
+
+    template<class TEvType>
+    struct TTestEventObserverTraits<TEventHandle<TEvType>>
+        : public TTestEventObserverTraits<TEvType>
+    {};
+
     class TTestActorRuntimeBase: public TNonCopyable {
     public:
         class TEdgeActor;
@@ -375,24 +405,19 @@ namespace NActors {
             observerHolder.Remove();
         */
 
-        template <typename TEvType>
+        template <typename TEvType = IEventHandle>
         TEventObserverHolder AddObserver(std::function<void(typename TEvType::TPtr&)> observerFunc)
         {
-            auto baseFunc = [observerFunc](TAutoPtr<IEventHandle>& event) {
-                if (event && event->GetTypeRewrite() == TEvType::EventType)
-                    observerFunc(*(reinterpret_cast<typename TEvType::TPtr*>(&event)));
+            auto baseFunc = [observerFunc](IEventHandle::TPtr& event) {
+                if (event && TTestEventObserverTraits<TEvType>::Match(event)) {
+                    observerFunc(TTestEventObserverTraits<TEvType>::Convert(event));
+                }
             };
 
             auto iter = ObserverFuncs.insert(ObserverFuncs.end(), baseFunc);
             return TEventObserverHolder(&ObserverFuncs, std::move(iter));
         }
 
-        TEventObserverHolder AddObserver(std::function<void(TAutoPtr<IEventHandle>&)> observerFunc)
-        {
-            auto iter = ObserverFuncs.insert(ObserverFuncs.end(), observerFunc);
-            return TEventObserverHolder(&ObserverFuncs, std::move(iter));
-        }
-
         template<typename T>
         void AppendToLogSettings(NLog::EComponent minVal, NLog::EComponent maxVal, T func) {
             Y_ABORT_UNLESS(!IsInitialized);
@@ -445,15 +470,14 @@ namespace NActors {
                 TDuration simTimeout = TDuration::Max())
         {
             typename TEvent::TPtr handle;
-            const ui32 eventType = TEvent::EventType;
             WaitForEdgeEvents([&](TTestActorRuntimeBase& runtime, TAutoPtr<IEventHandle>& event) {
                 Y_UNUSED(runtime);
-                if (event->GetTypeRewrite() != eventType)
+                if (!TTestEventObserverTraits<TEvent>::Match(event))
                     return false;
 
-                typename TEvent::TPtr* typedEvent = reinterpret_cast<typename TEvent::TPtr*>(&event);
-                if (predicate(*typedEvent)) {
-                    handle = *typedEvent;
+                typename TEvent::TPtr& typedEvent = TTestEventObserverTraits<TEvent>::Convert(event);
+                if (predicate(typedEvent)) {
+                    handle = std::move(typedEvent);
                     return true;
                 }
 
@@ -809,8 +833,8 @@ namespace NActors {
         const std::function<bool(const typename TEvent::TPtr&)>& predicate) {
         ev.Destroy();
         for (auto& event : events) {
-            if (event && event->GetTypeRewrite() == TEvent::EventType) {
-                if (predicate(reinterpret_cast<const typename TEvent::TPtr&>(event))) {
+            if (event && TTestEventObserverTraits<TEvent>::Match(event)) {
+                if (predicate(TTestEventObserverTraits<TEvent>::Convert(event))) {
                     ev = event;
                     return ev->CastAsLocal<TEvent>();
                 }