From 5481ab5ce5efc1a318e91895787acac94cbec229 Mon Sep 17 00:00:00 2001
From: Roger Johansson <roger@asynkron.se>
Date: Thu, 21 Aug 2025 21:35:36 +0200
Subject: [PATCH] Use TestProbe in timer tests

---
 .../Proto.Actor.Tests/TimerExtensionsTests.cs | 36 ++++++-------------
 1 file changed, 10 insertions(+), 26 deletions(-)

diff --git a/tests/Proto.Actor.Tests/TimerExtensionsTests.cs b/tests/Proto.Actor.Tests/TimerExtensionsTests.cs
index 35206472e6..a65f7e0a7b 100644
--- a/tests/Proto.Actor.Tests/TimerExtensionsTests.cs
+++ b/tests/Proto.Actor.Tests/TimerExtensionsTests.cs
@@ -1,5 +1,6 @@
 using System;
 using System.Threading.Tasks;
+using Proto.TestKit;
 using Proto.Timers;
 using Xunit;
 using Microsoft.Extensions.Time.Testing;
@@ -13,27 +14,18 @@ public async Task SchedulerSchedulesMessageAfterDelay()
     {
         await using var system = new ActorSystem();
         var context = system.Root;
-        var tcs = new TaskCompletionSource();
-
-        var pid = context.Spawn(Props.FromFunc(ctx =>
-        {
-            if (ctx.Message is "Wakeup")
-            {
-                tcs.SetResult();
-            }
-
-            return Task.CompletedTask;
-        }));
+        var probe = new TestProbe();
+        var pid = context.Spawn(Props.FromProducer(() => probe));
 
         var scheduler = context.Scheduler();
 
         scheduler.SendOnce(TimeSpan.FromMilliseconds(200), pid, "Wakeup");
 
         // ensure message isn't delivered immediately
-        await Task.Delay(100);
-        Assert.False(tcs.Task.IsCompleted);
+        await probe.ExpectNoMessageAsync(TimeSpan.FromMilliseconds(100));
 
-        await tcs.Task.WaitAsync(TimeSpan.FromSeconds(5));
+        var msg = await probe.GetNextMessageAsync<string>(TimeSpan.FromSeconds(5));
+        Assert.Equal("Wakeup", msg);
     }
 
     [Fact]
@@ -41,17 +33,8 @@ public async Task SchedulerWithTimeProviderSchedulesMessageAfterDelay()
     {
         await using var system = new ActorSystem();
         var context = system.Root;
-        var tcs = new TaskCompletionSource();
-
-        var pid = context.Spawn(Props.FromFunc(ctx =>
-        {
-            if (ctx.Message is "Wakeup")
-            {
-                tcs.SetResult();
-            }
-
-            return Task.CompletedTask;
-        }));
+        var probe = new TestProbe();
+        var pid = context.Spawn(Props.FromProducer(() => probe));
 
         var timeProvider = new FakeTimeProvider();
         var scheduler = context.Scheduler(timeProvider);
@@ -62,7 +45,8 @@ public async Task SchedulerWithTimeProviderSchedulesMessageAfterDelay()
         await Task.Delay(50);
         timeProvider.Advance(TimeSpan.FromMinutes(1));
 
-        await tcs.Task.WaitAsync(TimeSpan.FromMilliseconds(10));
+        var msg = await probe.GetNextMessageAsync<string>(TimeSpan.FromMilliseconds(10));
+        Assert.Equal("Wakeup", msg);
     }
 }