@@ -9,8 +9,6 @@ namespace Bunit.Extensions.WaitForHelpers;
99/// </summary>
1010public abstract class WaitForHelper < T > : IDisposable
1111{
12- private readonly object lockObject = new ( ) ;
13- private readonly Timer timer ;
1412 private readonly TaskCompletionSource < T > checkPassedCompletionSource ;
1513 private readonly Func < ( bool CheckPassed , T Content ) > completeChecker ;
1614 private readonly IRenderedFragmentBase renderedFragment ;
@@ -40,143 +38,154 @@ public abstract class WaitForHelper<T> : IDisposable
4038 /// </summary>
4139 public Task < T > WaitTask { get ; }
4240
43-
4441 /// <summary>
4542 /// Initializes a new instance of the <see cref="WaitForHelper{T}"/> class.
4643 /// </summary>
47- [ System . Diagnostics . CodeAnalysis . SuppressMessage ( "Performance" , "CA1849:Call async methods when in an async method" , Justification = "Using x.Result inside a ContinueWith is safe." ) ]
48- protected WaitForHelper ( IRenderedFragmentBase renderedFragment , Func < ( bool CheckPassed , T Content ) > completeChecker , TimeSpan ? timeout = null )
44+ protected WaitForHelper (
45+ IRenderedFragmentBase renderedFragment ,
46+ Func < ( bool CheckPassed , T Content ) > completeChecker ,
47+ TimeSpan ? timeout = null )
4948 {
5049 this . renderedFragment = renderedFragment ?? throw new ArgumentNullException ( nameof ( renderedFragment ) ) ;
5150 this . completeChecker = completeChecker ?? throw new ArgumentNullException ( nameof ( completeChecker ) ) ;
51+
5252 logger = renderedFragment . Services . CreateLogger < WaitForHelper < T > > ( ) ;
53+ checkPassedCompletionSource = new TaskCompletionSource < T > ( ) ;
54+ WaitTask = CreateWaitTask ( renderedFragment , timeout ) ;
5355
54- var renderer = renderedFragment . Services . GetRequiredService < ITestRenderer > ( ) ;
55- var renderException = renderer
56- . UnhandledException
57- . ContinueWith ( x => Task . FromException < T > ( x . Result ) , CancellationToken . None , TaskContinuationOptions . OnlyOnRanToCompletion , TaskScheduler . Current )
58- . Unwrap ( ) ;
56+ InitializeWaiting ( ) ;
57+ }
5958
60- checkPassedCompletionSource = new TaskCompletionSource < T > ( ) ;
61- WaitTask = Task . WhenAny ( checkPassedCompletionSource . Task , renderException ) . Unwrap ( ) ;
59+ /// <summary>
60+ /// Disposes the wait helper and cancels the any ongoing waiting, if it is not
61+ /// already in one of the other completed states.
62+ /// </summary>
63+ public void Dispose ( )
64+ {
65+ Dispose ( disposing : true ) ;
66+ GC . SuppressFinalize ( this ) ;
67+ }
6268
63- timer = new Timer ( OnTimeout , this , Timeout . InfiniteTimeSpan , Timeout . InfiniteTimeSpan ) ;
69+ /// <summary>
70+ /// Disposes of the wait task and related logic.
71+ /// </summary>
72+ /// <remarks>
73+ /// The disposing parameter should be false when called from a finalizer, and true when called from the
74+ /// <see cref="Dispose()"/> method. In other words, it is true when deterministically called and false when non-deterministically called.
75+ /// </remarks>
76+ /// <param name="disposing">Set to true if called from <see cref="Dispose()"/>, false if called from a finalizer.f.</param>
77+ protected virtual void Dispose ( bool disposing )
78+ {
79+ if ( isDisposed || ! disposing )
80+ return ;
81+
82+ isDisposed = true ;
83+ checkPassedCompletionSource . TrySetCanceled ( ) ;
84+ renderedFragment . OnAfterRender -= OnAfterRender ;
85+ logger . LogWaiterDisposed ( renderedFragment . ComponentId ) ;
86+ }
6487
88+ private void InitializeWaiting ( )
89+ {
6590 if ( ! WaitTask . IsCompleted )
6691 {
92+ var renderCountAtSubscribeTime = renderedFragment . RenderCount ;
93+
94+ // Before subscribing to renderedFragment.OnAfterRender,
95+ // we need to make sure that the desired state has not already been reached.
6796 OnAfterRender ( this , EventArgs . Empty ) ;
68- this . renderedFragment . OnAfterRender += OnAfterRender ;
69- OnAfterRender ( this , EventArgs . Empty ) ;
70- StartTimer ( timeout ) ;
97+
98+ SubscribeToOnAfterRender ( ) ;
99+
100+ // If the render count from before subscribing has changes
101+ // till now, we need to do trigger another check, since
102+ // the render may have happened asynchronously and before
103+ // the subscription was set up.
104+ if ( renderCountAtSubscribeTime < renderedFragment . RenderCount )
105+ {
106+ OnAfterRender ( this , EventArgs . Empty ) ;
107+ }
71108 }
72109 }
73110
74- private void StartTimer ( TimeSpan ? timeout )
111+ private Task < T > CreateWaitTask ( IRenderedFragmentBase renderedFragment , TimeSpan ? timeout )
75112 {
76- if ( isDisposed )
77- return ;
113+ var renderer = renderedFragment . Services . GetRequiredService < ITestRenderer > ( ) ;
78114
79- lock ( lockObject )
115+ // Two to failure conditions, that the renderer captures an unhandled
116+ // exception from a component or itself, or that the timeout is reached,
117+ // are executed on the renderes schedular, to ensure that OnAfterRender
118+ // and the continuations does not happen at the same time.
119+ var failureTask = renderer . Dispatcher . InvokeAsync ( ( ) =>
80120 {
81- if ( isDisposed )
82- return ;
83-
84- timer . Change ( GetRuntimeTimeout ( timeout ) , Timeout . InfiniteTimeSpan ) ;
85- }
121+ var taskScheduler = TaskScheduler . FromCurrentSynchronizationContext ( ) ;
122+
123+ var renderException = renderer
124+ . UnhandledException
125+ . ContinueWith (
126+ x => Task . FromException < T > ( x . Result ) ,
127+ CancellationToken . None ,
128+ TaskContinuationOptions . OnlyOnRanToCompletion | TaskContinuationOptions . ExecuteSynchronously ,
129+ taskScheduler ) ;
130+
131+ var timeoutTask = Task . Delay ( GetRuntimeTimeout ( timeout ) )
132+ . ContinueWith (
133+ x =>
134+ {
135+ logger . LogWaiterTimedOut ( renderedFragment . ComponentId ) ;
136+ return Task . FromException < T > ( new WaitForFailedException ( TimeoutErrorMessage , capturedException ) ) ;
137+ } ,
138+ CancellationToken . None ,
139+ TaskContinuationOptions . OnlyOnRanToCompletion | TaskContinuationOptions . ExecuteSynchronously ,
140+ taskScheduler ) ;
141+
142+ return Task . WhenAny ( renderException , timeoutTask ) . Unwrap ( ) ;
143+ } ) . Unwrap ( ) ;
144+
145+ return Task . WhenAny ( failureTask , checkPassedCompletionSource . Task ) . Unwrap ( ) ;
86146 }
87147
88148 private void OnAfterRender ( object ? sender , EventArgs args )
89149 {
90- if ( isDisposed )
150+ if ( isDisposed || WaitTask . IsCompleted )
91151 return ;
92152
93- lock ( lockObject )
153+ try
94154 {
95- if ( isDisposed )
96- return ;
155+ logger . LogCheckingWaitCondition ( renderedFragment . ComponentId ) ;
97156
98- try
157+ var checkResult = completeChecker ( ) ;
158+ if ( checkResult . CheckPassed )
99159 {
100- logger . LogCheckingWaitCondition ( renderedFragment . ComponentId ) ;
101-
102- var checkResult = completeChecker ( ) ;
103- if ( checkResult . CheckPassed )
104- {
105- checkPassedCompletionSource . TrySetResult ( checkResult . Content ) ;
106- logger . LogCheckCompleted ( renderedFragment . ComponentId ) ;
107- Dispose ( ) ;
108- }
109- else
110- {
111- logger . LogCheckFailed ( renderedFragment . ComponentId ) ;
112- }
160+ checkPassedCompletionSource . TrySetResult ( checkResult . Content ) ;
161+ logger . LogCheckCompleted ( renderedFragment . ComponentId ) ;
162+ Dispose ( ) ;
113163 }
114- catch ( Exception ex )
164+ else
115165 {
116- capturedException = ex ;
117- logger . LogCheckThrow ( renderedFragment . ComponentId , ex ) ;
118-
119- if ( StopWaitingOnCheckException )
120- {
121- checkPassedCompletionSource . TrySetException ( new WaitForFailedException ( CheckThrowErrorMessage , capturedException ) ) ;
122- Dispose ( ) ;
123- }
166+ logger . LogCheckFailed ( renderedFragment . ComponentId ) ;
124167 }
125168 }
126- }
127-
128- private void OnTimeout ( object ? state )
129- {
130- if ( isDisposed )
131- return ;
132-
133- lock ( lockObject )
169+ catch ( Exception ex )
134170 {
135- if ( isDisposed )
136- return ;
137-
138- logger . LogWaiterTimedOut ( renderedFragment . ComponentId ) ;
171+ capturedException = ex ;
172+ logger . LogCheckThrow ( renderedFragment . ComponentId , ex ) ;
139173
140- checkPassedCompletionSource . TrySetException ( new WaitForFailedException ( TimeoutErrorMessage , capturedException ) ) ;
141-
142- Dispose ( ) ;
174+ if ( StopWaitingOnCheckException )
175+ {
176+ checkPassedCompletionSource . TrySetException ( new WaitForFailedException ( CheckThrowErrorMessage , capturedException ) ) ;
177+ Dispose ( ) ;
178+ }
143179 }
144180 }
145181
146- /// <summary>
147- /// Disposes the wait helper and cancels the any ongoing waiting, if it is not
148- /// already in one of the other completed states.
149- /// </summary>
150- public void Dispose ( )
151- {
152- Dispose ( disposing : true ) ;
153- GC . SuppressFinalize ( this ) ;
154- }
155-
156- /// <summary>
157- /// Disposes of the wait task and related logic.
158- /// </summary>
159- /// <remarks>
160- /// The disposing parameter should be false when called from a finalizer, and true when called from the
161- /// <see cref="Dispose()"/> method. In other words, it is true when deterministically called and false when non-deterministically called.
162- /// </remarks>
163- /// <param name="disposing">Set to true if called from <see cref="Dispose()"/>, false if called from a finalizer.f.</param>
164- protected virtual void Dispose ( bool disposing )
182+ private void SubscribeToOnAfterRender ( )
165183 {
166- if ( isDisposed || ! disposing )
167- return ;
168-
169- lock ( lockObject )
170- {
171- if ( isDisposed )
172- return ;
173-
174- isDisposed = true ;
175- renderedFragment . OnAfterRender -= OnAfterRender ;
176- timer . Dispose ( ) ;
177- checkPassedCompletionSource . TrySetCanceled ( ) ;
178- logger . LogWaiterDisposed ( renderedFragment . ComponentId ) ;
179- }
184+ // There might not be a need to subscribe if the WaitTask has already
185+ // been completed, perhaps due to an unhandled exception from the
186+ // renderer or from the initial check by the checker.
187+ if ( ! isDisposed && ! WaitTask . IsCompleted )
188+ renderedFragment . OnAfterRender += OnAfterRender ;
180189 }
181190
182191 private static TimeSpan GetRuntimeTimeout ( TimeSpan ? timeout )
0 commit comments