@@ -203,6 +203,36 @@ namespace NActors {
203203 }
204204 };
205205
206+ /* *
207+ * Allows customizing behavior based on the event type
208+ */
209+ template <class TEvType >
210+ struct TTestEventObserverTraits {
211+ static bool Match (IEventHandle::TPtr& ev) noexcept {
212+ return ev->GetTypeRewrite () == TEvType::EventType;
213+ }
214+
215+ static typename TEvType::TPtr& Convert (IEventHandle::TPtr& ev) noexcept {
216+ return reinterpret_cast <typename TEvType::TPtr&>(ev);
217+ }
218+ };
219+
220+ template <>
221+ struct TTestEventObserverTraits <IEventHandle> {
222+ static constexpr bool Match (IEventHandle::TPtr&) noexcept {
223+ return true ;
224+ }
225+
226+ static constexpr IEventHandle::TPtr& Convert (IEventHandle::TPtr& ev) noexcept {
227+ return ev;
228+ }
229+ };
230+
231+ template <class TEvType >
232+ struct TTestEventObserverTraits <TEventHandle<TEvType>>
233+ : public TTestEventObserverTraits<TEvType>
234+ {};
235+
206236 class TTestActorRuntimeBase : public TNonCopyable {
207237 public:
208238 class TEdgeActor ;
@@ -375,24 +405,19 @@ namespace NActors {
375405 observerHolder.Remove();
376406 */
377407
378- template <typename TEvType>
408+ template <typename TEvType = IEventHandle >
379409 TEventObserverHolder AddObserver (std::function<void (typename TEvType::TPtr&)> observerFunc)
380410 {
381- auto baseFunc = [observerFunc](TAutoPtr<IEventHandle>& event) {
382- if (event && event->GetTypeRewrite () == TEvType::EventType)
383- observerFunc (*(reinterpret_cast <typename TEvType::TPtr*>(&event)));
411+ auto baseFunc = [observerFunc](IEventHandle::TPtr& event) {
412+ if (event && TTestEventObserverTraits<TEvType>::Match (event)) {
413+ observerFunc (TTestEventObserverTraits<TEvType>::Convert (event));
414+ }
384415 };
385416
386417 auto iter = ObserverFuncs.insert (ObserverFuncs.end (), baseFunc);
387418 return TEventObserverHolder (&ObserverFuncs, std::move (iter));
388419 }
389420
390- TEventObserverHolder AddObserver (std::function<void (TAutoPtr<IEventHandle>&)> observerFunc)
391- {
392- auto iter = ObserverFuncs.insert (ObserverFuncs.end (), observerFunc);
393- return TEventObserverHolder (&ObserverFuncs, std::move (iter));
394- }
395-
396421 template <typename T>
397422 void AppendToLogSettings (NLog::EComponent minVal, NLog::EComponent maxVal, T func) {
398423 Y_ABORT_UNLESS (!IsInitialized);
@@ -445,15 +470,14 @@ namespace NActors {
445470 TDuration simTimeout = TDuration::Max())
446471 {
447472 typename TEvent::TPtr handle;
448- const ui32 eventType = TEvent::EventType;
449473 WaitForEdgeEvents ([&](TTestActorRuntimeBase& runtime, TAutoPtr<IEventHandle>& event) {
450474 Y_UNUSED (runtime);
451- if (event-> GetTypeRewrite () != eventType )
475+ if (!TTestEventObserverTraits<TEvent>:: Match (event) )
452476 return false ;
453477
454- typename TEvent::TPtr* typedEvent = reinterpret_cast < typename TEvent::TPtr*>(& event);
455- if (predicate (* typedEvent)) {
456- handle = * typedEvent;
478+ typename TEvent::TPtr& typedEvent = TTestEventObserverTraits< TEvent>:: Convert ( event);
479+ if (predicate (typedEvent)) {
480+ handle = std::move ( typedEvent) ;
457481 return true ;
458482 }
459483
@@ -809,8 +833,8 @@ namespace NActors {
809833 const std::function<bool (const typename TEvent::TPtr&)>& predicate) {
810834 ev.Destroy ();
811835 for (auto & event : events) {
812- if (event && event-> GetTypeRewrite () == TEvent::EventType ) {
813- if (predicate (reinterpret_cast < const typename TEvent::TPtr&> (event))) {
836+ if (event && TTestEventObserverTraits< TEvent>:: Match (event) ) {
837+ if (predicate (TTestEventObserverTraits< TEvent>:: Convert (event))) {
814838 ev = event;
815839 return ev->CastAsLocal <TEvent>();
816840 }
0 commit comments