Add common WaitFor helper to the test actor runtime (#7725)

Author: snaury

ydb/core/testlib/actors/test_runtime.h

@@ -84,22 +84,42 @@ namespace NActors {
         void SimulateSleep(TDuration duration);
 
         template<class TResult>
-        inline TResult WaitFuture(NThreading::TFuture<TResult> f) {
+        inline TResult WaitFuture(NThreading::TFuture<TResult> f, TDuration simTimeout = TDuration::Max()) {
             if (!f.HasValue() && !f.HasException()) {
                 TDispatchOptions options;
                 options.CustomFinalCondition = [&]() {
                     return f.HasValue() || f.HasException();
                 };
-                options.FinalEvents.emplace_back([&](IEventHandle&) {
-                    return f.HasValue() || f.HasException();
-                });
+                // Quirk: non-empty FinalEvents enables full simulation
+                options.FinalEvents.emplace_back([](IEventHandle&) { return false; });
 
-                this->DispatchEvents(options);
+                this->DispatchEvents(options, simTimeout);
 
                 Y_ABORT_UNLESS(f.HasValue() || f.HasException());
             }
 
-            return f.ExtractValue();
+            if constexpr (!std::is_same_v<TResult, void>) {
+                return f.ExtractValue();
+            } else {
+                return f.GetValue();
+            }
+        }
+
+        template<class TCondition>
+        inline void WaitFor(const TString& description, const TCondition& condition, TDuration simTimeout = TDuration::Max()) {
+            if (!condition()) {
+                TDispatchOptions options;
+                options.CustomFinalCondition = [&]() {
+                    return condition();
+                };
+                // Quirk: non-empty FinalEvents enables full simulation
+                options.FinalEvents.emplace_back([](IEventHandle&) { return false; });
+
+                Cerr << "... waiting for " << description << Endl;
+                this->DispatchEvents(options, simTimeout);
+
+                Y_ABORT_UNLESS(condition(), "Timeout while waiting for %s", description.c_str());
+            }
         }
 
         void SendToPipe(ui64 tabletId, const TActorId& sender, IEventBase* payload, ui32 nodeIndex = 0,

ydb/core/testlib/actors/test_runtime_ut.cpp

@@ -622,6 +622,101 @@ Y_UNIT_TEST_SUITE(TActorTest) {
             UNIT_ASSERT_VALUES_EQUAL(event->Get()->Index, 12u);
         }
     }
+
+    Y_UNIT_TEST(TestWaitFuture) {
+        enum EEv {
+            EvTrigger = EventSpaceBegin(TEvents::ES_PRIVATE)
+        };
+
+        struct TEvTrigger : public TEventLocal<TEvTrigger, EvTrigger> {
+            TEvTrigger() = default;
+        };
+
+        class TTriggerActor : public TActorBootstrapped<TTriggerActor> {
+        public:
+            TTriggerActor(NThreading::TPromise<void> promise)
+                : Promise(std::move(promise))
+            {}
+
+            void Bootstrap() {
+                Schedule(TDuration::Seconds(1), new TEvTrigger);
+                Become(&TThis::StateWork);
+            }
+
+        private:
+            STFUNC(StateWork) {
+                switch (ev->GetTypeRewrite()) {
+                    hFunc(TEvTrigger, Handle);
+                }
+            }
+
+            void Handle(TEvTrigger::TPtr&) {
+                Promise.SetValue();
+                PassAway();
+            }
+
+        private:
+            NThreading::TPromise<void> Promise;
+        };
+
+        TTestActorRuntime runtime;
+        runtime.Initialize(MakeEgg());
+
+        NThreading::TPromise<void> promise = NThreading::NewPromise<void>();
+        NThreading::TFuture<void> future = promise.GetFuture();
+
+        auto actor = runtime.Register(new TTriggerActor(std::move(promise)));
+        runtime.EnableScheduleForActor(actor);
+
+        runtime.WaitFuture(std::move(future));
+    }
+
+    Y_UNIT_TEST(TestWaitFor) {
+        enum EEv {
+            EvTrigger = EventSpaceBegin(TEvents::ES_PRIVATE)
+        };
+
+        struct TEvTrigger : public TEventLocal<TEvTrigger, EvTrigger> {
+            TEvTrigger() = default;
+        };
+
+        class TTriggerActor : public TActorBootstrapped<TTriggerActor> {
+        public:
+            TTriggerActor(int* ptr)
+                : Ptr(ptr)
+            {}
+
+            void Bootstrap() {
+                Schedule(TDuration::Seconds(1), new TEvTrigger);
+                Become(&TThis::StateWork);
+            }
+
+        private:
+            STFUNC(StateWork) {
+                switch (ev->GetTypeRewrite()) {
+                    hFunc(TEvTrigger, Handle);
+                }
+            }
+
+            void Handle(TEvTrigger::TPtr&) {
+                *Ptr = 42;
+                PassAway();
+            }
+
+        private:
+            int* Ptr;
+        };
+
+        TTestActorRuntime runtime;
+        runtime.Initialize(MakeEgg());
+
+        int value = 0;
+        auto actor = runtime.Register(new TTriggerActor(&value));
+        runtime.EnableScheduleForActor(actor);
+
+        runtime.WaitFor("value = 42", [&]{ return value == 42; });
+        UNIT_ASSERT_VALUES_EQUAL(value, 42);
+    }
 };
 
 }