feat: Add async overloads

Author: linkdotnet

CHANGELOG.md

@@ -6,6 +6,10 @@ All notable changes to **bUnit** will be documented in this file. The project ad
 
 ## [Unreleased]
 
+### Added
+
+- New overloads to WaitForHelpers to have async assertions and predicates. Reported by [@radmorecameron](https://github.com/radmorecameron) in #1833. Fixed by [@linkdotnet](https://github.com/linkdotnet).
+
 ## [2.7.2] - 2026-03-31
 
 ### Fixed

src/bunit/Extensions/WaitForHelpers/RenderedComponentWaitForHelperExtensions.WaitForState.cs

@@ -56,6 +56,31 @@ public static async Task WaitForStateAsync<TComponent>(this IRenderedComponent<T
 			await waiter.WaitTask;
 		}
 
+		/// <summary>
+		/// Wait until the provided asynchronous <paramref name="statePredicate"/> returns true,
+		/// or the <paramref name="timeout"/> is reached (default is one second).
+		///
+		/// The <paramref name="statePredicate"/> is evaluated initially, and then each time
+		/// the <paramref name="renderedComponent"/> renders.
+		/// </summary>
+		/// <param name="renderedComponent">The render fragment or component to attempt to verify state against.</param>
+		/// <param name="statePredicate">The asynchronous predicate to invoke after each render, which must return <c>true</c> when the desired state has been reached.</param>
+		/// <param name="timeout">The maximum time to wait for the desired state.</param>
+		/// <exception cref="WaitForFailedException">Thrown if the <paramref name="statePredicate"/> throw an exception during invocation, or if the timeout has been reached. See the inner exception for details.</exception>
+		/// <remarks>
+		/// The predicate is awaited on the renderer's dispatcher. Because awaiting yields the dispatcher,
+		/// a render may interleave at each <c>await</c> point, so the predicate is not guaranteed an atomic
+		/// view of the DOM across awaits. Use this overload when the check itself must <c>await</c> (for
+		/// example an async service or JavaScript interop call); prefer the synchronous <see cref="WaitForState{TComponent}(IRenderedComponent{TComponent}, Func{bool}, TimeSpan?)"/>
+		/// overload for pure DOM or markup checks.
+		/// </remarks>
+		public static async Task WaitForStateAsync<TComponent>(this IRenderedComponent<TComponent> renderedComponent, Func<Task<bool>> statePredicate, TimeSpan? timeout = null)
+			where TComponent : IComponent
+		{
+			using var waiter = new WaitForStateHelper<TComponent>(renderedComponent, statePredicate, timeout);
+			await waiter.WaitTask;
+		}
+
 		/// <summary>
 		/// Wait until the provided <paramref name="assertion"/> passes (i.e. does not throw an
 		/// exception), or the <paramref name="timeout"/> is reached (default is one second).
@@ -99,4 +124,29 @@ public static async Task WaitForAssertionAsync<TComponent>(this IRenderedCompone
 			using var waiter = new WaitForAssertionHelper<TComponent>(renderedComponent, assertion, timeout);
 			await waiter.WaitTask;
 		}
+
+		/// <summary>
+		/// Wait until the provided asynchronous <paramref name="assertion"/> passes (i.e. does not throw an
+		/// exception), or the <paramref name="timeout"/> is reached (default is one second).
+		///
+		/// The <paramref name="assertion"/> is attempted initially, and then each time the <paramref name="renderedComponent"/> renders.
+		/// </summary>
+		/// <param name="renderedComponent">The rendered fragment to wait for renders from and assert against.</param>
+		/// <param name="assertion">The asynchronous verification or assertion to perform.</param>
+		/// <param name="timeout">The maximum time to attempt the verification.</param>
+		/// <exception cref="WaitForFailedException">Thrown if the timeout has been reached. See the inner exception to see the captured assertion exception.</exception>
+		/// <remarks>
+		/// The assertion is awaited on the renderer's dispatcher. Because awaiting yields the dispatcher,
+		/// a render may interleave at each <c>await</c> point, so the assertion is not guaranteed an atomic
+		/// view of the DOM across awaits. Use this overload when the assertion itself must <c>await</c> (for
+		/// example an async service or JavaScript interop call); prefer the synchronous <see cref="WaitForAssertion{TComponent}(IRenderedComponent{TComponent}, Action, TimeSpan?)"/>
+		/// overload for pure DOM or markup assertions.
+		/// </remarks>
+		[AssertionMethod]
+		public static async Task WaitForAssertionAsync<TComponent>(this IRenderedComponent<TComponent> renderedComponent, Func<Task> assertion, TimeSpan? timeout = null)
+			where TComponent : IComponent
+		{
+			using var waiter = new WaitForAssertionHelper<TComponent>(renderedComponent, assertion, timeout);
+			await waiter.WaitTask;
+		}
 }

src/bunit/Extensions/WaitForHelpers/WaitForAssertionHelper.cs

@@ -37,4 +37,28 @@ public WaitForAssertionHelper(IRenderedComponent<TComponent> renderedComponent,
 			  },
 			  timeout)
 	{ }
+
+	/// <summary>
+	/// Initializes a new instance of the <see cref="WaitForAssertionHelper{TComponent}"/> class,
+	/// which will until the provided asynchronous <paramref name="assertion"/> passes (i.e. does not throw an
+	/// exception), or the <paramref name="timeout"/> is reached (default is one second).
+	///
+	/// The <paramref name="assertion"/> is attempted initially, and then each time the <paramref name="renderedComponent"/> renders.
+	/// </summary>
+	/// <param name="renderedComponent">The rendered fragment to wait for renders from and assert against.</param>
+	/// <param name="assertion">The asynchronous verification or assertion to perform.</param>
+	/// <param name="timeout">The maximum time to attempt the verification.</param>
+	/// <remarks>
+	/// If a debugger is attached the timeout is set to <see cref="Timeout.InfiniteTimeSpan" />, giving the possibility to debug without the timeout triggering.
+	/// </remarks>
+	public WaitForAssertionHelper(IRenderedComponent<TComponent> renderedComponent, Func<Task> assertion, TimeSpan? timeout = null)
+		: base(
+			  renderedComponent,
+			  async () =>
+			  {
+				  await assertion();
+				  return (true, default);
+			  },
+			  timeout)
+	{ }
 }

src/bunit/Extensions/WaitForHelpers/WaitForHelper.cs

@@ -12,12 +12,14 @@ internal abstract class WaitForHelper<T, TComponent> : IDisposable
 	where TComponent : IComponent
 {
 	private readonly TaskCompletionSource<T> checkPassedCompletionSource;
-	private readonly Func<(bool CheckPassed, T Content)> completeChecker;
+	private readonly Func<ValueTask<(bool CheckPassed, T Content)>> completeChecker;
 	private readonly IRenderedComponent<TComponent> renderedComponent;
 	private readonly ILogger<WaitForHelper<T, TComponent>> logger;
 	private readonly BunitRenderer renderer;
 	private readonly Timer? timer;
 	private bool isDisposed;
+	private bool isChecking;
+	private bool isDirty;
 	private int checkCount;
 	private Exception? capturedException;
 
@@ -44,12 +46,31 @@ internal abstract class WaitForHelper<T, TComponent> : IDisposable
 	public Task<T> WaitTask { get; }
 
 	/// <summary>
-	/// Initializes a new instance of the <see cref="WaitForHelper{T, TComponent}"/> class.
+	/// Initializes a new instance of the <see cref="WaitForHelper{T, TComponent}"/> class
+	/// with a synchronous check. The check is wrapped into a completed <see cref="ValueTask{TResult}"/>,
+	/// so it is still evaluated synchronously (with no render interleaving) on the renderer's dispatcher.
 	/// </summary>
 	protected WaitForHelper(
 		IRenderedComponent<TComponent> renderedComponent,
 		Func<(bool CheckPassed, T Content)> completeChecker,
 		TimeSpan? timeout = null)
+		: this(renderedComponent, WrapSynchronousChecker(completeChecker), timeout)
+	{
+	}
+
+	/// <summary>
+	/// Initializes a new instance of the <see cref="WaitForHelper{T, TComponent}"/> class.
+	/// </summary>
+	/// <remarks>
+	/// The <paramref name="completeChecker"/> is awaited on the renderer's dispatcher after each render.
+	/// A synchronous check (see the other constructor) wraps into an already-completed <see cref="ValueTask{TResult}"/>
+	/// and therefore resumes synchronously without yielding the dispatcher; a genuinely asynchronous check yields the
+	/// dispatcher at each <c>await</c>, so a render may interleave and the check is not guaranteed an atomic view of the DOM.
+	/// </remarks>
+	protected WaitForHelper(
+		IRenderedComponent<TComponent> renderedComponent,
+		Func<ValueTask<(bool CheckPassed, T Content)>> completeChecker,
+		TimeSpan? timeout = null)
 	{
 		this.renderedComponent = renderedComponent ?? throw new ArgumentNullException(nameof(renderedComponent));
 		this.completeChecker = completeChecker ?? throw new ArgumentNullException(nameof(completeChecker));
@@ -166,40 +187,87 @@ private void OnAfterRender(object? sender, EventArgs args)
 		if (isDisposed || WaitTask.IsCompleted)
 			return;
 
-		try
+		// The check may be asynchronous, so it cannot be awaited inside this synchronous
+		// event handler. Kick it off without blocking the dispatcher. A synchronous check
+		// completes inline (see RunCheckAsync). If a check is already in flight, mark the
+		// state dirty so the running check re-evaluates once it completes, ensuring the
+		// final render is never missed.
+		if (isChecking)
 		{
-			logger.LogCheckingWaitCondition(renderedComponent.ComponentId);
+			isDirty = true;
+			return;
+		}
 
-			var checkResult = completeChecker();
-			checkCount++;
-			if (checkResult.CheckPassed)
-			{
-				checkPassedCompletionSource.TrySetResult(checkResult.Content);
-				logger.LogCheckCompleted(renderedComponent.ComponentId);
-				Dispose();
-			}
-			else
+		isChecking = true;
+		_ = RunCheckAsync();
+	}
+
+	private async Task RunCheckAsync()
+	{
+		// Runs on the renderer's dispatcher. A synchronous check wraps into an already-completed
+		// ValueTask, so the await below resumes inline without yielding - preserving the atomic,
+		// no-render-interleaving behaviour synchronous checks rely on. A genuinely asynchronous
+		// check yields the dispatcher; the default await captures and resumes on the dispatcher's
+		// synchronization context, so isChecking/isDirty are only ever touched on the single
+		// dispatcher thread and need no locking.
+		try
+		{
+			do
 			{
-				logger.LogCheckFailed(renderedComponent.ComponentId);
+				isDirty = false;
+
+				if (isDisposed || WaitTask.IsCompleted)
+					return;
+
+				try
+				{
+					logger.LogCheckingWaitCondition(renderedComponent.ComponentId);
+
+					var checkResult = await completeChecker();
+					checkCount++;
+
+					if (isDisposed || WaitTask.IsCompleted)
+						return;
+
+					if (checkResult.CheckPassed)
+					{
+						checkPassedCompletionSource.TrySetResult(checkResult.Content);
+						logger.LogCheckCompleted(renderedComponent.ComponentId);
+						Dispose();
+						return;
+					}
+
+					logger.LogCheckFailed(renderedComponent.ComponentId);
+				}
+				catch (Exception ex)
+				{
+					checkCount++;
+					capturedException = ex;
+					logger.LogCheckThrow(renderedComponent.ComponentId, ex);
+
+					if (StopWaitingOnCheckException)
+					{
+						if (!isDisposed && !WaitTask.IsCompleted)
+						{
+							checkPassedCompletionSource.TrySetException(
+								new WaitForFailedException(
+									CheckThrowErrorMessage ?? string.Empty,
+									checkCount,
+									renderedComponent.RenderCount,
+									renderer.RenderCount,
+									capturedException));
+							Dispose();
+						}
+
+						return;
+					}
+				}
 			}
+			while (isDirty && !isDisposed && !WaitTask.IsCompleted);
 		}
-		catch (Exception ex)
+		finally
 		{
-			checkCount++;
-			capturedException = ex;
-			logger.LogCheckThrow(renderedComponent.ComponentId, ex);
-
-			if (StopWaitingOnCheckException)
-			{
-				checkPassedCompletionSource.TrySetException(
-					new WaitForFailedException(
-						CheckThrowErrorMessage ?? string.Empty,
-						checkCount,
-						renderedComponent.RenderCount,
-						renderer.RenderCount,
-						capturedException));
-				Dispose();
-			}
+			isChecking = false;
 		}
 	}
 
@@ -218,4 +286,14 @@ private static TimeSpan GetRuntimeTimeout(TimeSpan? timeout)
 			? Timeout.InfiniteTimeSpan
 			: timeout ?? BunitContext.DefaultWaitTimeout;
 	}
+
+	private static Func<ValueTask<(bool CheckPassed, T Content)>> WrapSynchronousChecker(Func<(bool CheckPassed, T Content)> completeChecker)
+	{
+		ArgumentNullException.ThrowIfNull(completeChecker);
+
+		// Wrap the synchronous check in an already-completed ValueTask so the single check
+		// pipeline can await it. Awaiting a completed ValueTask resumes inline, so the check
+		// still runs synchronously on the dispatcher with no render interleaving.
+		return () => new ValueTask<(bool CheckPassed, T Content)>(completeChecker());
+	}
 }

src/bunit/Extensions/WaitForHelpers/WaitForStateHelper.cs

@@ -34,4 +34,21 @@ public WaitForStateHelper(IRenderedComponent<TComponent> renderedComponent, Func
 		: base(renderedComponent, () => (statePredicate(), default), timeout)
 	{
 	}
+
+	/// <summary>
+	/// Initializes a new instance of the <see cref="WaitForStateHelper{TComponent}"/> class,
+	/// which will wait until the provided asynchronous <paramref name="statePredicate"/> returns true,
+	/// or the <paramref name="timeout"/> is reached (default is one second).
+	/// </summary>
+	/// <remarks>
+	/// The <paramref name="statePredicate"/> is evaluated initially, and then each time the <paramref name="renderedComponent"/> renders.
+	/// </remarks>
+	/// <param name="renderedComponent">The render fragment or component to attempt to verify state against.</param>
+	/// <param name="statePredicate">The asynchronous predicate to invoke after each render, which must return <c>true</c> when the desired state has been reached.</param>
+	/// <param name="timeout">The maximum time to wait for the desired state.</param>
+	/// <exception cref="WaitForFailedException">Thrown if the <paramref name="statePredicate"/> throw an exception during invocation, or if the timeout has been reached. See the inner exception for details.</exception>
+	public WaitForStateHelper(IRenderedComponent<TComponent> renderedComponent, Func<Task<bool>> statePredicate, TimeSpan? timeout = null)
+		: base(renderedComponent, async () => (await statePredicate(), default), timeout)
+	{
+	}
 }