Handle synchronous TestFailure within async fn body

lib/src/dom/async/wait_for.dart

@@ -71,6 +71,7 @@ Future<T> waitFor<T>(
   MutationObserver observer;
   Timer intervalTimer;
   Timer overallTimeoutTimer;
+  Stopwatch stopwatch;
   bool resultPending;
   final doneCompleter = Completer<T>();
 
@@ -115,7 +116,25 @@ Future<T> waitFor<T>(
         resultPending = true;
         await (result as Future).then((resolvedValue) {
           onDone(null, resolvedValue as T);
-        }, onError: onDone).timeout(timeout, onTimeout: handleTimeout);
+        }).catchError((error) {
+          onDone(error, result);
+        }, test: (error) {
+          if (error is TestFailure) {
+            // There was an test `expect()` within an async function body that threw,
+            // not an *actual* future that resulted in an unexpected error. In this case,
+            // keep trying until time is up by restarting the `overallTimeoutTimer` that was
+            // cancelled using the time remaining.
+            overallTimeoutTimer = Timer(timeout - stopwatch.elapsed, handleTimeout);
+            // Save the most recent test failure so we can throw it in the event of a timeout.
+            lastError = error;
+            // Do not call `catchError` since we want to keep retrying.
+            return false;
+          }
+
+          // The error was not a TestFailure coming from an `expect()` within an async
+          // function body, so return true to cause `catchError` to be called.
+          return true;
+        }).timeout(timeout, onTimeout: handleTimeout);
       } else {
         onDone(null, result);
       }
@@ -127,6 +146,7 @@ Future<T> waitFor<T>(
   }
 
   overallTimeoutTimer = Timer(timeout, handleTimeout);
+  stopwatch = Stopwatch()..start();
 
   intervalTimer = Timer.periodic(interval, checkCallback);
   observer = MutationObserver((_, __) => checkCallback())

test/unit/dom/wait_for_test.dart

@@ -58,6 +58,30 @@ void main() {
                     rtl.waitFor(() => expect(view.container.contains(rootElement), isFalse), container: view.container),
                 throwsA(isA<TestFailure>()));
           }, timeout: asyncQueryTestTimeout);
+
+          group('that is placed within an asynchronous function body', () {
+            test('and succeeds before timeout', () async {
+              var numRuns = 0;
+              await rtl.waitFor(() async {
+                numRuns++;
+                expect(numRuns, 5);
+              });
+            }, timeout: asyncQueryTestTimeout);
+
+            test(
+                'and never succeeds before timeout, throwing the consumer TestFailure instead of a generic TimeoutException',
+                () async {
+              var numRuns = 0;
+              expect(
+                  () => rtl.waitFor(() async {
+                        if (numRuns < 4) {
+                          numRuns++;
+                        }
+                        expect(numRuns, 5);
+                      }),
+                  throwsA(isA<TestFailure>()));
+            }, timeout: asyncQueryTestTimeout);
+          });
         });
 
         test('a function that throws an arbitrary error, rethrows the error thrown by the expectation', () async {