Test runtime: Easy wait for a single event (#8608)

Author: azevaykin

ydb/core/testlib/actors/block_events.h

@@ -86,5 +86,4 @@ namespace NActors {
         bool Stopped = false;
     };
 
-
 } // namespace NActors

ydb/core/testlib/actors/test_runtime_ut.cpp

@@ -1,5 +1,6 @@
 #include <ydb/core/testlib/actors/test_runtime.h>
 #include <ydb/core/testlib/actors/block_events.h>
+#include <ydb/core/testlib/actors/wait_events.h>
 #include <ydb/core/base/appdata.h>
 #include <ydb/library/actors/core/event_local.h>
 #include <ydb/library/actors/core/events.h>
@@ -826,6 +827,65 @@ Y_UNIT_TEST_SUITE(TActorTest) {
         UNIT_ASSERT_VALUES_EQUAL(values.at(1), 4);
         UNIT_ASSERT_VALUES_EQUAL(values.at(2), 7);
     }
+
+    Y_UNIT_TEST(TestWaitForFirstEvent) {
+        enum EEv {
+            EvTrigger = EventSpaceBegin(TEvents::ES_PRIVATE)
+        };
+
+        struct TEvTrigger : public TEventLocal<TEvTrigger, EvTrigger> {
+            int Value;
+
+            TEvTrigger(int value)
+                : Value(value)
+            {}
+        };
+
+        class TSourceActor : public TActorBootstrapped<TSourceActor> {
+        public:
+            TSourceActor(const TActorId& target)
+                : Target(target)
+            {}
+
+            void Bootstrap() {
+                Become(&TThis::StateWork);
+                Schedule(TDuration::Seconds(1), new TEvents::TEvWakeup);
+            }
+
+        private:
+            STFUNC(StateWork) {
+                switch (ev->GetTypeRewrite()) {
+                    hFunc(TEvents::TEvWakeup, Handle);
+                }
+            }
+
+            void Handle(TEvents::TEvWakeup::TPtr&) {
+                Send(Target, new TEvTrigger(++Counter));
+                Schedule(TDuration::Seconds(1), new TEvents::TEvWakeup);
+            }
+
+        private:
+            TActorId Target;
+            int Counter = 0;
+        };
+
+        TTestActorRuntime runtime;
+        runtime.Initialize(MakeEgg());
+
+        TActorId target = runtime.AllocateEdgeActor();
+
+        TActorId actorId = runtime.Register(new TSourceActor(target));
+        runtime.EnableScheduleForActor(actorId);
+
+        {
+            TWaitForFirstEvent<TEvTrigger> waiter(runtime);
+            waiter.Wait();
+        }
+        {
+            TWaitForFirstEvent<TEvTrigger> waiter(runtime, [](const TEvTrigger::TPtr& ev){ return ev->Get()->Value == 10; });
+            waiter.Wait();
+        }
+    }    
 }
 
 }

ydb/core/testlib/actors/wait_events.cpp

@@ -0,0 +1 @@
+#include "wait_events.h"

ydb/core/testlib/actors/wait_events.h

@@ -0,0 +1,48 @@
+#include "test_runtime.h"
+
+#include <functional>
+
+namespace NActors {
+
+    /**
+     * Easy wait for a first event under the test actor runtime
+     *
+     */
+    template<class TEvType>
+    class TWaitForFirstEvent {
+    public:
+        TWaitForFirstEvent(TTestActorRuntime& runtime, std::function<bool(const typename TEvType::TPtr&)> condition = {})
+            : Runtime(runtime)
+            , Condition(std::move(condition))
+            , Holder(Runtime.AddObserver<TEvType>(
+                [this](typename TEvType::TPtr& ev) {
+                    if (EventSeen)
+                        return;
+                    if (Condition && !Condition(ev))
+                        return;
+                    EventSeen = true;
+                }))
+        {}
+
+        /**
+         * Wait for a first event
+         */
+        void Wait() {
+            Runtime.WaitFor(TypeName<TEvType>(), [&]{ return EventSeen; });
+        }
+
+        /**
+         * Stops waiting and remove event observer
+         */
+        void Stop() {
+            Holder.Remove();
+        }
+
+    private:
+        TTestActorRuntime& Runtime;
+        std::function<bool(const typename TEvType::TPtr&)> Condition;
+        TTestActorRuntime::TEventObserverHolder Holder;
+        bool EventSeen = false;
+    };
+
+} // namespace NActors

ydb/core/testlib/actors/ya.make

@@ -5,6 +5,8 @@ SRCS(
     block_events.h
     test_runtime.cpp
     test_runtime.h
+    wait_events.cpp
+    wait_events.h
 )
 
 PEERDIR(