Changed WaitFor helper methods to throw any unhandled exceptions caug…

…ht by the the renderer while waiting for state or assertion

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

@@ -1,4 +1,5 @@
 using System;
+using System.Runtime.ExceptionServices;
 using Bunit.Extensions.WaitForHelpers;
 
 namespace Bunit
@@ -22,13 +23,21 @@ public static class RenderedFragmentWaitForHelperExtensions
 		public static void WaitForState(this IRenderedFragmentBase renderedFragment, Func<bool> statePredicate, TimeSpan? timeout = null)
 		{
 			using var waiter = new WaitForStateHelper(renderedFragment, statePredicate, timeout);
+
 			try
 			{
 				waiter.WaitTask.GetAwaiter().GetResult();
 			}
-			catch (AggregateException e) when (e.InnerException is not null)
+			catch (Exception e)
 			{
-				throw e.InnerException;
+				if (e is AggregateException aggregateException && aggregateException.InnerExceptions.Count == 1)
+				{
+					ExceptionDispatchInfo.Capture(aggregateException.InnerExceptions[0]).Throw();
+				}
+				else
+				{
+					ExceptionDispatchInfo.Capture(e).Throw();
+				}
 			}
 		}
 
@@ -45,13 +54,21 @@ public static void WaitForState(this IRenderedFragmentBase renderedFragment, Fun
 		public static void WaitForAssertion(this IRenderedFragmentBase renderedFragment, Action assertion, TimeSpan? timeout = null)
 		{
 			using var waiter = new WaitForAssertionHelper(renderedFragment, assertion, timeout);
+
 			try
 			{
 				waiter.WaitTask.GetAwaiter().GetResult();
 			}
-			catch (AggregateException e) when (e.InnerException is not null)
+			catch (Exception e)
 			{
-				throw e.InnerException;
+				if (e is AggregateException aggregateException && aggregateException.InnerExceptions.Count == 1)
+				{
+					ExceptionDispatchInfo.Capture(aggregateException.InnerExceptions[0]).Throw();
+				}
+				else
+				{
+					ExceptionDispatchInfo.Capture(e).Throw();
+				}
 			}
 		}
 	}

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

@@ -1,6 +1,8 @@
 using System;
 using System.Threading;
 using System.Threading.Tasks;
+using Bunit.Rendering;
+using Microsoft.Extensions.DependencyInjection;
 using Microsoft.Extensions.Logging;
 
 namespace Bunit.Extensions.WaitForHelpers
@@ -13,7 +15,7 @@ public abstract class WaitForHelper : IDisposable
 	{
 		private readonly object lockObject = new();
 		private readonly Timer timer;
-		private readonly TaskCompletionSource<object?> completionSouce;
+		private readonly TaskCompletionSource<object?> checkPassedCompletionSouce;
 		private readonly Func<bool> completeChecker;
 		private readonly IRenderedFragmentBase renderedFragment;
 		private readonly ILogger logger;
@@ -40,7 +42,7 @@ public abstract class WaitForHelper : IDisposable
 		/// Gets the task that will complete successfully if the check passed before the timeout was reached.
 		/// The task will complete with an <see cref="WaitForFailedException"/> exception if the timeout was reached without the check passing.
 		/// </summary>
-		public Task WaitTask => completionSouce.Task;
+		public Task WaitTask { get; }
 
 		/// <summary>
 		/// Initializes a new instance of the <see cref="WaitForHelper"/> class.
@@ -50,13 +52,25 @@ protected WaitForHelper(IRenderedFragmentBase renderedFragment, Func<bool> compl
 			this.renderedFragment = renderedFragment ?? throw new ArgumentNullException(nameof(renderedFragment));
 			this.completeChecker = completeChecker ?? throw new ArgumentNullException(nameof(completeChecker));
 			logger = renderedFragment.Services.CreateLogger<WaitForHelper>();
-			completionSouce = new TaskCompletionSource<object?>();
+
+			var renderer = renderedFragment.Services.GetRequiredService<ITestRenderer>();
+			var renderException = renderer
+				.UnhandledException
+				.ContinueWith(x => Task.FromException(x.Result), CancellationToken.None, TaskContinuationOptions.OnlyOnRanToCompletion, TaskScheduler.Current)
+				.Unwrap();
+
+			checkPassedCompletionSouce = new TaskCompletionSource<object?>();
+			WaitTask = Task.WhenAny(checkPassedCompletionSouce.Task, renderException).Unwrap();
+
 			timer = new Timer(OnTimeout, this, Timeout.InfiniteTimeSpan, Timeout.InfiniteTimeSpan);
 
-			OnAfterRender(this, EventArgs.Empty);
-			this.renderedFragment.OnAfterRender += OnAfterRender;
-			OnAfterRender(this, EventArgs.Empty);
-			StartTimer(timeout);
+			if (!WaitTask.IsCompleted)
+			{
+				OnAfterRender(this, EventArgs.Empty);
+				this.renderedFragment.OnAfterRender += OnAfterRender;
+				OnAfterRender(this, EventArgs.Empty);
+				StartTimer(timeout);
+			}
 		}
 
 		private void StartTimer(TimeSpan? timeout)
@@ -88,7 +102,7 @@ private void OnAfterRender(object? sender, EventArgs args)
 					logger.LogDebug(new EventId(1, nameof(OnAfterRender)), $"Checking the wait condition for component {renderedFragment.ComponentId}");
 					if (completeChecker())
 					{
-						completionSouce.TrySetResult(null);
+						checkPassedCompletionSouce.TrySetResult(null);
 						logger.LogDebug(new EventId(2, nameof(OnAfterRender)), $"The check completed successfully for component {renderedFragment.ComponentId}");
 						Dispose();
 					}
@@ -104,7 +118,7 @@ private void OnAfterRender(object? sender, EventArgs args)
 
 					if (StopWaitingOnCheckException)
 					{
-						completionSouce.TrySetException(new WaitForFailedException(CheckThrowErrorMessage, capturedException));
+						checkPassedCompletionSouce.TrySetException(new WaitForFailedException(CheckThrowErrorMessage, capturedException));
 						Dispose();
 					}
 				}
@@ -123,7 +137,7 @@ private void OnTimeout(object? state)
 
 				logger.LogDebug(new EventId(5, nameof(OnTimeout)), $"The wait for helper for component {renderedFragment.ComponentId} timed out");
 
-				completionSouce.TrySetException(new WaitForFailedException(TimeoutErrorMessage, capturedException));
+				checkPassedCompletionSouce.TrySetException(new WaitForFailedException(TimeoutErrorMessage, capturedException));
 
 				Dispose();
 			}
@@ -160,7 +174,7 @@ protected virtual void Dispose(bool disposing)
 				isDisposed = true;
 				renderedFragment.OnAfterRender -= OnAfterRender;
 				timer.Dispose();
-				completionSouce.TrySetCanceled();
+				checkPassedCompletionSouce.TrySetCanceled();
 				logger.LogDebug(new EventId(6, nameof(Dispose)), $"The state wait helper for component {renderedFragment.ComponentId} disposed");
 			}
 		}