fix: Wait for dispatcher before starting the timer

Author: linkdotnet

src/bunit/Extensions/WaitForHelpers/WaitForHelper.cs

@@ -69,28 +69,45 @@ protected WaitForHelper(
 		// resulting in the timer completing before a single check has
 		// has a chance to complete.
 		WaitTask = CreateWaitTask();
-		CheckAndInitializeWaiting();
+		var initializationTask = CheckAndInitializeWaiting();
 
 		// If the initial check did not complete successfully,
 		// start the timer and recheck after every render until the timer expires.
+		// Wait for the initialization to complete to ensure the OnAfterRender
+		// subscription is registered before the timer can expire.
+		// This prevents a race condition where the timer fires before renders
+		// are being monitored, causing WaitForState to timeout prematurely.
 		if (!WaitTask.IsCompleted)
 		{
-			timer = new Timer(
-				static (state) =>
-				{
-					var @this = (WaitForHelper<T, TComponent>)state!;
-					@this.logger.LogWaiterTimedOut(@this.renderedComponent.ComponentId);
-					@this.checkPassedCompletionSource.TrySetException(
-						new WaitForFailedException(
-							@this.TimeoutErrorMessage ?? string.Empty,
-							@this.checkCount,
-							@this.renderedComponent.RenderCount,
-							@this.renderer.RenderCount,
-							@this.capturedException));
-				},
-				this,
-				GetRuntimeTimeout(timeout),
-				Timeout.InfiniteTimeSpan);
+			// Ensure subscription is complete before starting the timer.
+			// On slower systems, the InvokeAsync queue may be delayed,
+			// and we need to ensure OnAfterRender is subscribed before
+			// the timeout timer starts ticking.
+			if (!initializationTask.IsCompleted)
+			{
+				initializationTask.GetAwaiter().GetResult();
+			}
+
+			// Only start the timer if we still haven't completed after initialization
+			if (!WaitTask.IsCompleted)
+			{
+				timer = new Timer(
+					static (state) =>
+					{
+						var @this = (WaitForHelper<T, TComponent>)state!;
+						@this.logger.LogWaiterTimedOut(@this.renderedComponent.ComponentId);
+						@this.checkPassedCompletionSource.TrySetException(
+							new WaitForFailedException(
+								@this.TimeoutErrorMessage ?? string.Empty,
+								@this.checkCount,
+								@this.renderedComponent.RenderCount,
+								@this.renderer.RenderCount,
+								@this.capturedException));
+					},
+					this,
+					GetRuntimeTimeout(timeout),
+					Timeout.InfiniteTimeSpan);
+			}
 		}
 	}
 
@@ -124,7 +141,7 @@ protected virtual void Dispose(bool disposing)
 		logger.LogWaiterDisposed(renderedComponent.ComponentId);
 	}
 
-	private void CheckAndInitializeWaiting()
+	private Task CheckAndInitializeWaiting()
 	{
 		if (!WaitTask.IsCompleted)
 		{
@@ -134,14 +151,16 @@ private void CheckAndInitializeWaiting()
 			// This also ensures that checks performed during OnAfterRender,
 			// which are usually not atomic, e.g. search the DOM tree,
 			// can be performed without the DOM tree changing.
-			renderedComponent.InvokeAsync(() =>
+			return renderedComponent.InvokeAsync(() =>
 			{
 				// Before subscribing to renderedFragment.OnAfterRender,
 				// we need to make sure that the desired state has not already been reached.
 				OnAfterRender(this, EventArgs.Empty);
 				SubscribeToOnAfterRender();
 			});
 		}
+
+		return Task.CompletedTask;
 	}
 
 	private Task<T> CreateWaitTask()