From ac958b90f9daea860d71cce789af3532bcfcc095 Mon Sep 17 00:00:00 2001
From: Ricky Kaare Engelharth <ricky@ricky.dk>
Date: Wed, 2 Oct 2024 08:10:58 +0200
Subject: [PATCH] Make WaitForCall extension accept matching calls made before
 waiting

---
 src/Atc.Test/SubstituteExtensions.cs | 207 ++++++++++++++++++---------
 1 file changed, 137 insertions(+), 70 deletions(-)

diff --git a/src/Atc.Test/SubstituteExtensions.cs b/src/Atc.Test/SubstituteExtensions.cs
index 593680c..83dddd7 100644
--- a/src/Atc.Test/SubstituteExtensions.cs
+++ b/src/Atc.Test/SubstituteExtensions.cs
@@ -1,4 +1,6 @@
 // ReSharper disable AsyncVoidLambda
+using NSubstitute.Exceptions;
+
 namespace Atc.Test;
 
 /// <summary>
@@ -83,20 +85,29 @@ public static async Task WaitForCall<T>(
         TimeSpan timeout = default)
         where T : class
     {
-        var completion = new TaskCompletionSource<bool>();
-        substitute
-            .When(substituteCall)
-            .Do(_ => completion.TrySetResult(true));
+        try
+        {
+            substitute
+                .ValidateCallReceived(
+                    substituteCall,
+                    MatchArgs.AsSpecifiedInCall);
+        }
+        catch (ReceivedCallsException)
+        {
+            var completion = new TaskCompletionSource<bool>();
+            substitute
+                .When(substituteCall)
+                .Do(_ => completion.TrySetResult(true));
 
-        await completion
-            .WaitForCompletion(timeout)
-            .ConfigureAwait(false);
+            await completion
+                .WaitForCompletion(timeout)
+                .ConfigureAwait(false);
 
-        substitute
-            .ValidateCallReceived(
-                substituteCall
-                ?? throw new ArgumentNullException(nameof(substituteCall)),
-                MatchArgs.AsSpecifiedInCall);
+            substitute
+                .ValidateCallReceived(
+                    substituteCall,
+                    MatchArgs.AsSpecifiedInCall);
+        }
     }
 
     /// <summary>
@@ -114,19 +125,29 @@ public static async Task WaitForCall<T>(
         TimeSpan timeout = default)
         where T : class
     {
-        var completion = new TaskCompletionSource<bool>();
-        substitute
-            .When(substituteCall)
-            .Do(_ => completion.TrySetResult(true));
+        try
+        {
+            substitute
+                .ValidateCallReceived(
+                    x => substituteCall(x),
+                    MatchArgs.AsSpecifiedInCall);
+        }
+        catch (ReceivedCallsException)
+        {
+            var completion = new TaskCompletionSource<bool>();
+            substitute
+                .When(substituteCall)
+                .Do(_ => completion.TrySetResult(true));
 
-        await completion
-            .WaitForCompletion(timeout)
-            .ConfigureAwait(false);
+            await completion
+                .WaitForCompletion(timeout)
+                .ConfigureAwait(false);
 
-        substitute
-            .ValidateCallReceived(
-                x => substituteCall(x),
-                MatchArgs.AsSpecifiedInCall);
+            substitute
+                .ValidateCallReceived(
+                    x => substituteCall(x),
+                    MatchArgs.AsSpecifiedInCall);
+        }
     }
 
     /// <summary>
@@ -139,26 +160,40 @@ await completion
     /// <param name="substituteCall">The call to wait for.</param>
     /// <param name="timeout">Timeout for the wait operation.</param>
     /// <returns>A task representing the async operation.</returns>
+    [SuppressMessage(
+        "Reliability",
+        "CA2012:Use ValueTasks correctly",
+        Justification = "Call should not be awaited for route to work")]
     public static async Task WaitForCall<TSubstitute, TResult>(
         this TSubstitute substitute,
         Func<TSubstitute, ValueTask<TResult>> substituteCall,
         TimeSpan timeout = default)
         where TSubstitute : class
     {
-        var completion = new TaskCompletionSource<bool>();
-        substitute
-            .When(substituteCall)
-            .Do(_ => completion.TrySetResult(true));
+        try
+        {
+            substitute
+                .ValidateCallReceived(
+                    x => substituteCall(x),
+                    MatchArgs.AsSpecifiedInCall);
+        }
+        catch (ReceivedCallsException)
+        {
+            var completion = new TaskCompletionSource<bool>();
+            substitute
+                .When(substituteCall)
+                .Do(_ => completion.TrySetResult(true));
 
-        await completion
-            .WaitForCompletion(timeout)
-            .ConfigureAwait(false);
+            await completion
+                .WaitForCompletion(timeout)
+                .ConfigureAwait(false);
 
-        substitute
-            .ValidateCallReceived(
-                async x => await substituteCall(x)
-                    .ConfigureAwait(false),
-                MatchArgs.AsSpecifiedInCall);
+            substitute
+                .ValidateCallReceived(
+                    x => substituteCall(x)
+                        .ConfigureAwait(false),
+                    MatchArgs.AsSpecifiedInCall);
+        }
     }
 
     /// <summary>
@@ -176,20 +211,29 @@ public static async Task WaitForCallForAnyArgs<T>(
         TimeSpan timeout = default)
         where T : class
     {
-        var completion = new TaskCompletionSource<bool>();
-        substitute
-            .WhenForAnyArgs(substituteCall)
-            .Do(_ => completion.TrySetResult(true));
+        try
+        {
+            substitute
+                .ValidateCallReceived(
+                    substituteCall,
+                    MatchArgs.Any);
+        }
+        catch (ReceivedCallsException)
+        {
+            var completion = new TaskCompletionSource<bool>();
+            substitute
+                .WhenForAnyArgs(substituteCall)
+                .Do(_ => completion.TrySetResult(true));
 
-        await completion
-            .WaitForCompletion(timeout)
-            .ConfigureAwait(false);
+            await completion
+                .WaitForCompletion(timeout)
+                .ConfigureAwait(false);
 
-        substitute
-            .ValidateCallReceived(
-                substituteCall
-                ?? throw new ArgumentNullException(nameof(substituteCall)),
-                MatchArgs.Any);
+            substitute
+                .ValidateCallReceived(
+                    substituteCall,
+                    MatchArgs.Any);
+        }
     }
 
     /// <summary>
@@ -207,19 +251,29 @@ public static async Task WaitForCallForAnyArgs<T>(
         TimeSpan timeout = default)
         where T : class
     {
-        var completion = new TaskCompletionSource<bool>();
-        substitute
-            .WhenForAnyArgs(substituteCall)
-            .Do(_ => completion.TrySetResult(true));
+        try
+        {
+            substitute
+                .ValidateCallReceived(
+                    x => substituteCall(x),
+                    MatchArgs.Any);
+        }
+        catch (ReceivedCallsException)
+        {
+            var completion = new TaskCompletionSource<bool>();
+            substitute
+                .WhenForAnyArgs(substituteCall)
+                .Do(_ => completion.TrySetResult(true));
 
-        await completion
-            .WaitForCompletion(timeout)
-            .ConfigureAwait(false);
+            await completion
+                .WaitForCompletion(timeout)
+                .ConfigureAwait(false);
 
-        substitute
-            .ValidateCallReceived(
-                x => substituteCall(x),
-                MatchArgs.Any);
+            substitute
+                .ValidateCallReceived(
+                    x => substituteCall(x),
+                    MatchArgs.Any);
+        }
     }
 
     /// <summary>
@@ -232,26 +286,39 @@ await completion
     /// <param name="substituteCall">The call to wait for.</param>
     /// <param name="timeout">Timeout for the wait operation.</param>
     /// <returns>A task representing the async operation.</returns>
+    [SuppressMessage(
+        "Reliability",
+        "CA2012:Use ValueTasks correctly",
+        Justification = "Call should not be awaited for route to work")]
     public static async Task WaitForCallForAnyArgs<TSubstitute, TResult>(
         this TSubstitute substitute,
         Func<TSubstitute, ValueTask<TResult>> substituteCall,
         TimeSpan timeout = default)
         where TSubstitute : class
     {
-        var completion = new TaskCompletionSource<bool>();
-        substitute
-            .WhenForAnyArgs(substituteCall)
-            .Do(_ => completion.TrySetResult(true));
+        try
+        {
+            substitute
+                .ValidateCallReceived(
+                    x => substituteCall(x),
+                    MatchArgs.Any);
+        }
+        catch (ReceivedCallsException)
+        {
+            var completion = new TaskCompletionSource<bool>();
+            substitute
+                .WhenForAnyArgs(substituteCall)
+                .Do(_ => completion.TrySetResult(true));
 
-        await completion
-            .WaitForCompletion(timeout)
-            .ConfigureAwait(false);
+            await completion
+                .WaitForCompletion(timeout)
+                .ConfigureAwait(false);
 
-        substitute
-            .ValidateCallReceived(
-                async x => await substituteCall(x)
-                    .ConfigureAwait(false),
-                MatchArgs.Any);
+            substitute
+                .ValidateCallReceived(
+                    x => substituteCall(x),
+                    MatchArgs.Any);
+        }
     }
 
     private static void ValidateCallReceived<T>(