Extend test/task cancellation support to test case evaluation.

Sources/Testing/Issues/Issue+Recording.swift

@@ -186,8 +186,7 @@ extension Issue {
       // This error is thrown by expectation checking functions to indicate a
       // condition evaluated to `false`. Those functions record their own issue,
       // so we don't need to record another one redundantly.
-    } catch is SkipInfo,
-            is CancellationError where Task.isCancelled {
+    } catch let error where SkipInfo(error) != nil {
       // This error represents control flow rather than an issue, so we suppress
       // it here.
     } catch {
@@ -232,8 +231,7 @@ extension Issue {
       // This error is thrown by expectation checking functions to indicate a
       // condition evaluated to `false`. Those functions record their own issue,
       // so we don't need to record another one redundantly.
-    } catch is SkipInfo,
-            is CancellationError where Task.isCancelled {
+    } catch let error where SkipInfo(error) != nil {
       // This error represents control flow rather than an issue, so we suppress
       // it here.
     } catch {

Sources/Testing/Running/Runner.Plan.swift

@@ -201,47 +201,75 @@ extension Runner.Plan {
   /// - Parameters:
   ///   - test: The test whose action will be determined.
   ///
-  /// - Returns: A tuple containing the action to take for `test` as well as any
-  ///   error that was thrown during trait evaluation. If more than one error
-  ///   was thrown, the first-caught error is returned.
-  private static func _determineAction(for test: Test) async -> (Action, (any Error)?) {
+  /// - Returns:The action to take for `test`.
+  private static func _determineAction(for test: inout Test) async -> Action {
+    let result: Action
+
     // We use a task group here with a single child task so that, if the trait
     // code calls Test.cancel() we don't end up cancelling the entire test run.
     // We could also model this as an unstructured task except that they aren't
     // available in the "task-to-thread" concurrency model.
     //
     // FIXME: Parallelize this work. Calling `prepare(...)` on all traits and
     // evaluating all test arguments should be safely parallelizable.
-    await withTaskGroup(returning: (Action, (any Error)?).self) { taskGroup in
+    (test, result) = await withTaskGroup(returning: (Test, Action).self) { [test] taskGroup in
       taskGroup.addTask {
+        var test = test
         var action = _runAction
-        var firstCaughtError: (any Error)?
 
         await Test.withCurrent(test) {
-          for trait in test.traits {
+          do {
+            var firstCaughtError: (any Error)?
+
+            for trait in test.traits {
+              do {
+                try await trait.prepare(for: test)
+              } catch {
+                if let skipInfo = SkipInfo(error) {
+                  action = .skip(skipInfo)
+                  break
+                } else {
+                  // Only preserve the first caught error
+                  firstCaughtError = firstCaughtError ?? error
+                }
+              }
+            }
+
+            // If no trait specified that the test should be skipped, but one
+            // did throw an error, then the action is to record an issue for
+            // that error.
+            if case .run = action, let error = firstCaughtError {
+              action = .recordIssue(Issue(for: error))
+            }
+          }
+
+          // If the test is still planned to run (i.e. nothing thus far has
+          // caused it to be skipped), evaluate its test cases now.
+          //
+          // The argument expressions of each test are captured in closures so
+          // they can be evaluated lazily only once it is determined that the
+          // test will run, to avoid unnecessary work. But now is the
+          // appropriate time to evaluate them.
+          if case .run = action {
             do {
-              try await trait.prepare(for: test)
-            } catch let error as SkipInfo {
-              action = .skip(error)
-              break
-            } catch is CancellationError where Task.isCancelled {
-              // Synthesize skip info for this cancellation error.
-              let sourceContext = SourceContext(backtrace: .current(), sourceLocation: nil)
-              let skipInfo = SkipInfo(comment: nil, sourceContext: sourceContext)
-              action = .skip(skipInfo)
-              break
+              try await test.evaluateTestCases()
             } catch {
-              // Only preserve the first caught error
-              firstCaughtError = firstCaughtError ?? error
+              if let skipInfo = SkipInfo(error) {
+                action = .skip(skipInfo)
+              } else {
+                action = .recordIssue(Issue(for: error))
+              }
             }
           }
         }
 
-        return (action, firstCaughtError)
+        return (test, action)
       }
 
       return await taskGroup.first { _ in true }!
     }
+
+    return result
   }
 
   /// Construct a graph of runner plan steps for the specified tests.
@@ -309,36 +337,12 @@ extension Runner.Plan {
         return nil
       }
 
-      var action = runAction
-      var firstCaughtError: (any Error)?
-
       // Walk all the traits and tell each to prepare to run the test.
       // If any throw a `SkipInfo` error at this stage, stop walking further.
       // But if any throw another kind of error, keep track of the first error
       // but continue walking, because if any subsequent traits throw a
       // `SkipInfo`, the error should not be recorded.
-      (action, firstCaughtError) = await _determineAction(for: test)
-
-      // If no trait specified that the test should be skipped, but one did
-      // throw an error, then the action is to record an issue for that error.
-      if case .run = action, let error = firstCaughtError {
-        action = .recordIssue(Issue(for: error))
-      }
-
-      // If the test is still planned to run (i.e. nothing thus far has caused
-      // it to be skipped), evaluate its test cases now.
-      //
-      // The argument expressions of each test are captured in closures so they
-      // can be evaluated lazily only once it is determined that the test will
-      // run, to avoid unnecessary work. But now is the appropriate time to
-      // evaluate them.
-      if case .run = action {
-        do {
-          try await test.evaluateTestCases()
-        } catch {
-          action = .recordIssue(Issue(for: error))
-        }
-      }
+      var action = await _determineAction(for: &test)
 
       // If the test is parameterized but has no cases, mark it as skipped.
       if case .run = action, let testCases = test.testCases, testCases.first(where: { _ in true }) == nil {

Sources/Testing/Running/SkipInfo.swift

@@ -54,6 +54,30 @@ extension SkipInfo: Equatable, Hashable {}
 
 extension SkipInfo: Codable {}
 
+// MARK: -
+
+extension SkipInfo {
+  /// Initialize an instance of this type from an arbitrary error.
+  ///
+  /// - Parameters:
+  ///   - error: The error to convert to an instance of this type.
+  ///
+  /// If `error` does not represent a skip or cancellation event, this
+  /// initializer returns `nil`.
+  init?(_ error: any Error) {
+    if let skipInfo = error as? Self {
+      self = skipInfo
+    } else if error is CancellationError, Task.isCancelled {
+      // Synthesize skip info for this cancellation error.
+      let backtrace = Backtrace(forFirstThrowOf: error) ?? .current()
+      let sourceContext = SourceContext(backtrace: backtrace, sourceLocation: nil)
+      self.init(comment: nil, sourceContext: sourceContext)
+    } else {
+      return nil
+    }
+  }
+}
+
 // MARK: - Deprecated
 
 extension SkipInfo {

Tests/TestingTests/TestCancellationTests.swift

@@ -132,6 +132,22 @@
     }
   }
 
+  @Test func `Cancelling a test while evaluating test cases skips the test`() async {
+    await testCancellation(testSkipped: 1) { configuration in
+      await Test(arguments: { try await cancelledTestCases(cancelsTask: false) }) { _ in
+        Issue.record("Recorded an issue!")
+      }.run(configuration: configuration)
+    }
+  }
+
+  @Test func `Cancelling the current task while evaluating test cases skips the test`() async {
+    await testCancellation(testSkipped: 1) { configuration in
+      await Test(arguments: { try await cancelledTestCases(cancelsTask: true) }) { _ in
+        Issue.record("Recorded an issue!")
+      }.run(configuration: configuration)
+    }
+  }
+
 #if !SWT_NO_EXIT_TESTS
   @Test func `Cancelling the current test from within an exit test`() async {
     await testCancellation(testCancelled: 1, testCaseCancelled: 1) { configuration in
@@ -219,6 +235,15 @@ struct CancelledTrait: TestTrait {
   }
 }
 
+func cancelledTestCases(cancelsTask: Bool) async throws -> EmptyCollection<Int> {
+  if cancelsTask {
+    withUnsafeCurrentTask { $0?.cancel() }
+    try Task.checkCancellation()
+  }
+  try Test.cancel("Cancelled from trait")
+}
+
+
 #if !SWT_NO_SNAPSHOT_TYPES
 struct `Shows as skipped in Xcode 16` {
   @Test func `Cancelled test`() throws {

Tests/TestingTests/TestSupport/TestingAdditions.swift

@@ -199,6 +199,23 @@ extension Test {
     self.init(name: name, displayName: name, traits: traits, sourceLocation: sourceLocation, containingTypeInfo: nil, testCases: caseGenerator, parameters: parameters)
   }
 
+  init<C>(
+    _ traits: any TestTrait...,
+    arguments collection: @escaping @Sendable () async throws -> C,
+    parameters: [Parameter] = [
+      Parameter(index: 0, firstName: "x", type: C.Element.self),
+    ],
+    sourceLocation: SourceLocation = #_sourceLocation,
+    column: Int = #column,
+    name: String = #function,
+    testFunction: @escaping @Sendable (C.Element) async throws -> Void
+  ) where C: Collection & Sendable, C.Element: Sendable {
+    let caseGenerator = { @Sendable in
+      Case.Generator(arguments: try await collection(), parameters: parameters, testFunction: testFunction)
+    }
+    self.init(name: name, displayName: name, traits: traits, sourceLocation: sourceLocation, containingTypeInfo: nil, testCases: caseGenerator, parameters: parameters)
+  }
+
   /// Initialize an instance of this type with a function or closure to call,
   /// parameterized over two collections of values.
   ///