use task.delay instead of timer

Author: egil

src/bunit.core/Extensions/WaitForHelpers/WaitForHelper.cs

@@ -9,8 +9,6 @@ namespace Bunit.Extensions.WaitForHelpers;
 /// </summary>
 public abstract class WaitForHelper<T> : IDisposable
 {
-	//// private readonly object lockObject = new();
-	private readonly Timer timer;
 	private readonly TaskCompletionSource<T> checkPassedCompletionSource;
 	private readonly Func<(bool CheckPassed, T Content)> completeChecker;
 	private readonly IRenderedFragmentBase renderedFragment;
@@ -40,60 +38,108 @@ public abstract class WaitForHelper<T> : IDisposable
 	/// </summary>
 	public Task<T> WaitTask { get; }
 
-
 	/// <summary>
 	/// Initializes a new instance of the <see cref="WaitForHelper{T}"/> class.
 	/// </summary>
-	protected WaitForHelper(IRenderedFragmentBase renderedFragment, Func<(bool CheckPassed, T Content)> completeChecker, TimeSpan? timeout = null)
+	protected WaitForHelper(
+		IRenderedFragmentBase renderedFragment,
+		Func<(bool CheckPassed, T Content)> completeChecker,
+		TimeSpan? timeout = null)
 	{
 		this.renderedFragment = renderedFragment ?? throw new ArgumentNullException(nameof(renderedFragment));
 		this.completeChecker = completeChecker ?? throw new ArgumentNullException(nameof(completeChecker));
 		logger = renderedFragment.Services.CreateLogger<WaitForHelper<T>>();
+		checkPassedCompletionSource = new TaskCompletionSource<T>();
 
-		var renderer = renderedFragment.Services.GetRequiredService<ITestRenderer>();
-		var renderException = renderer
-			.UnhandledException
-			.ContinueWith(x => Task.FromException<T>(x.Result), CancellationToken.None, TaskContinuationOptions.OnlyOnRanToCompletion, TaskScheduler.Current)
-			.Unwrap();
+		WaitTask = CreateWaitTask(renderedFragment, timeout);
 
-		checkPassedCompletionSource = new TaskCompletionSource<T>();
-		WaitTask = Task.WhenAny(checkPassedCompletionSource.Task, renderException).Unwrap();
+		InitializeWaiting();
+	}
+
+	/// <summary>
+	/// Disposes the wait helper and cancels the any ongoing waiting, if it is not
+	/// already in one of the other completed states.
+	/// </summary>
+	public void Dispose()
+	{
+		Dispose(disposing: true);
+		GC.SuppressFinalize(this);
+	}
+
+	/// <summary>
+	/// Disposes of the wait task and related logic.
+	/// </summary>
+	/// <remarks>
+	/// The disposing parameter should be false when called from a finalizer, and true when called from the
+	/// <see cref="Dispose()"/> method. In other words, it is true when deterministically called and false when non-deterministically called.
+	/// </remarks>
+	/// <param name="disposing">Set to true if called from <see cref="Dispose()"/>, false if called from a finalizer.f.</param>
+	protected virtual void Dispose(bool disposing)
+	{
+		if (isDisposed || !disposing)
+			return;
 
-		timer = new Timer(OnTimeout, this, Timeout.InfiniteTimeSpan, Timeout.InfiniteTimeSpan);
+		isDisposed = true;
+		checkPassedCompletionSource.TrySetCanceled();
+		renderedFragment.OnAfterRender -= OnAfterRender;
+		logger.LogWaiterDisposed(renderedFragment.ComponentId);
+	}
 
+	private void InitializeWaiting()
+	{
 		if (!WaitTask.IsCompleted)
 		{
+			var renderCountAtSubscribeTime = renderedFragment.RenderCount;
+
+			// Before subscribing to renderedFragment.OnAfterRender,
+			// we need to make sure that the desired state has not already been reached.
 			OnAfterRender(this, EventArgs.Empty);
-			this.renderedFragment.OnAfterRender += OnAfterRender;
-			OnAfterRender(this, EventArgs.Empty);
-			StartTimer(timeout);
+
+			SubscribeToOnAfterRender();
+
+			// If the render count from before subscribing has changes
+			// till now, we need to do trigger another check, since
+			// the render may have happened asynchronously and before
+			// the subscription was set up.
+			if (renderCountAtSubscribeTime < renderedFragment.RenderCount)
+			{
+				OnAfterRender(this, EventArgs.Empty);
+			}
 		}
 	}
 
-	private void StartTimer(TimeSpan? timeout)
+	private Task<T> CreateWaitTask(IRenderedFragmentBase renderedFragment, TimeSpan? timeout)
 	{
-		if (isDisposed)
-			return;
+		var renderer = renderedFragment.Services.GetRequiredService<ITestRenderer>();
 
-		////lock (lockObject)
-		////{
-		////	if (isDisposed)
-		////		return;
+		var renderException = renderer
+			.UnhandledException
+			.ContinueWith(
+				x => Task.FromException<T>(x.Result),
+				CancellationToken.None,
+				TaskContinuationOptions.OnlyOnRanToCompletion,
+				TaskScheduler.Current)
+			.Unwrap();
 
-		timer.Change(GetRuntimeTimeout(timeout), Timeout.InfiniteTimeSpan);
-		////}
+		var timeoutTask = Task.Delay(GetRuntimeTimeout(timeout))
+			.ContinueWith(
+				x => Task.FromException<T>(new WaitForFailedException(TimeoutErrorMessage, capturedException)),
+				CancellationToken.None,
+				TaskContinuationOptions.OnlyOnRanToCompletion,
+				TaskScheduler.Current)
+			.Unwrap();
+
+		return Task.WhenAny(
+			checkPassedCompletionSource.Task,
+			renderException,
+			timeoutTask).Unwrap();
 	}
 
 	private void OnAfterRender(object? sender, EventArgs args)
 	{
 		if (isDisposed)
 			return;
 
-		////lock (lockObject)
-		////{
-		////	if (isDisposed)
-		////		return;
-
 		try
 		{
 			logger.LogCheckingWaitCondition(renderedFragment.ComponentId);
@@ -121,61 +167,12 @@ private void OnAfterRender(object? sender, EventArgs args)
 				Dispose();
 			}
 		}
-		////}
-	}
-
-	private void OnTimeout(object? state)
-	{
-		if (isDisposed)
-			return;
-
-		////lock (lockObject)
-		////{
-		////	if (isDisposed)
-		////		return;
-
-		logger.LogWaiterTimedOut(renderedFragment.ComponentId);
-
-		checkPassedCompletionSource.TrySetException(new WaitForFailedException(TimeoutErrorMessage, capturedException));
-
-		Dispose();
-		////}
-	}
-
-	/// <summary>
-	/// Disposes the wait helper and cancels the any ongoing waiting, if it is not
-	/// already in one of the other completed states.
-	/// </summary>
-	public void Dispose()
-	{
-		Dispose(disposing: true);
-		GC.SuppressFinalize(this);
 	}
 
-	/// <summary>
-	/// Disposes of the wait task and related logic.
-	/// </summary>
-	/// <remarks>
-	/// The disposing parameter should be false when called from a finalizer, and true when called from the
-	/// <see cref="Dispose()"/> method. In other words, it is true when deterministically called and false when non-deterministically called.
-	/// </remarks>
-	/// <param name="disposing">Set to true if called from <see cref="Dispose()"/>, false if called from a finalizer.f.</param>
-	protected virtual void Dispose(bool disposing)
+	private void SubscribeToOnAfterRender()
 	{
-		if (isDisposed || !disposing)
-			return;
-
-		////lock (lockObject)
-		////{
-		////	if (isDisposed)
-		////		return;
-
-		isDisposed = true;
-		renderedFragment.OnAfterRender -= OnAfterRender;
-		timer.Dispose();
-		checkPassedCompletionSource.TrySetCanceled();
-		logger.LogWaiterDisposed(renderedFragment.ComponentId);
-		////}
+		if (!isDisposed && !WaitTask.IsCompleted)
+			renderedFragment.OnAfterRender += OnAfterRender;
 	}
 
 	private static TimeSpan GetRuntimeTimeout(TimeSpan? timeout)