Move test repetition down to the test case level

Sources/Testing/ABI/Encoded/ABI.EncodedEvent.swift

@@ -118,9 +118,14 @@ extension ABI {
         kind = .testStarted
       case .testCaseStarted:
         if eventContext.test?.isParameterized == false {
-          return nil
+          if let iteration = eventContext.iteration, iteration > 1 {
+            kind = .testStarted
+          } else {
+            return nil
+          }
+        } else {
+          kind = .testCaseStarted
         }
-        kind = .testCaseStarted
       case let .issueRecorded(recordedIssue):
         kind = .issueRecorded
         issue = EncodedIssue(encoding: recordedIssue, in: eventContext)
@@ -129,9 +134,14 @@ extension ABI {
         self.attachment = EncodedAttachment(encoding: attachment)
       case .testCaseEnded:
         if eventContext.test?.isParameterized == false {
-          return nil
+          if let iteration = eventContext.iteration, iteration > 1 {
+            kind = .testEnded
+          } else {
+            return nil
+          }
+        } else {
+          kind = .testCaseEnded
         }
-        kind = .testCaseEnded
       case .testCaseCancelled:
         kind = .testCaseCancelled
       case .testEnded:
@@ -164,6 +174,7 @@ extension ABI {
           let .testCancelled(skipInfo):
           _comments = Array(skipInfo.comment).map(\.rawValue)
           _sourceLocation = skipInfo.sourceLocation.map { EncodedSourceLocation(encoding: $0) }
+          _iteration = eventContext.iteration
         default:
           break
         }

Sources/Testing/Events/Event.swift

@@ -238,7 +238,6 @@ public struct Event: Sendable {
   static func post(
     _ kind: Kind,
     for testAndTestCase: (Test?, Test.Case?) = currentTestAndTestCase(),
-    iteration: Int? = nil,
     instant: Test.Clock.Instant = .now,
     configuration: Configuration? = nil
   ) {
@@ -262,7 +261,12 @@ public struct Event: Sendable {
       }
     }
     let event = Event(kind, testID: test?.id, testCaseID: testCase?.id, instant: instant)
-    let context = Event.Context(test: test, testCase: testCase, iteration: iteration, configuration: nil)
+    let context = Event.Context(
+      test: test,
+      testCase: testCase,
+      iteration: Test.currentIteration,
+      configuration: nil
+    )
     event._post(in: context, configuration: configuration)
   }
 }
@@ -327,8 +331,6 @@ extension Event {
       iteration: Int?,
       configuration: Configuration?
     ) {
-      // Ensure that if `iteration` is specified, the test is also specified.
-      precondition(iteration == nil || (iteration != nil && test != nil))
       self.test = test
       self.testCase = testCase
       self.iteration = iteration

Sources/Testing/Running/Configuration.swift

@@ -132,6 +132,10 @@ public struct Configuration: Sendable {
     }
   }
 
+  /// Whether to perform test repetition at the plan level or on a per-test-
+  /// case basis.
+  public var shouldUseLegacyPlanLevelRepetition: Bool = true
+
   /// Whether or not, and how, to iterate the test plan repeatedly.
   ///
   /// By default, the value of this property allows for a single iteration.

Sources/Testing/Running/Runner.RuntimeState.swift

@@ -29,6 +29,9 @@ extension Runner {
     /// The test case that is running on the current task, if any.
     var testCase: Test.Case?
 
+    /// The current iteration of the test repetition policy.
+    var iteration: Int?
+
     /// The runtime state related to the runner running on the current task,
     /// if any.
     @TaskLocal
@@ -233,6 +236,27 @@ extension Test {
       try await test.withCancellationHandling(body)
     }
   }
+
+  static var currentIteration: Int? {
+    Runner.RuntimeState.current?.iteration
+  }
+
+  /// Call a function while the value of ``Test/currentIteration`` is set.
+  ///
+  /// - Parameters:
+  ///   - iteration: The new value to set for ``Test/currentIteration``.
+  ///   - body: A function to call.
+  ///
+  /// - Returns: Whatever is returned by `body`.
+  ///
+  /// - Throws: Whatever is thrown by `body`.
+  static func withCurrentIteration<R>(_ iteration: Int?, perform body: () async throws -> R) async rethrows -> R {
+    var runtimeState = Runner.RuntimeState.current ?? .init()
+    runtimeState.iteration = iteration
+    return try await Runner.RuntimeState.$current.withValue(runtimeState) {
+      try await body()
+    }
+  }
 }
 
 extension Test.Case {

Sources/Testing/Running/Runner.swift

@@ -82,9 +82,6 @@ extension Runner {
   private struct _Context: Sendable {
     /// A serializer used to reduce parallelism among test cases.
     var testCaseSerializer: Serializer<Void>?
-
-    /// Which iteration of the test plan is being executed.
-    var iteration: Int
   }
 
   /// Apply the custom scope for any test scope providers of the traits
@@ -230,10 +227,10 @@ extension Runner {
       // Determine what kind of event to send for this step based on its action.
       switch step.action {
       case .run:
-        Event.post(.testStarted, for: (step.test, nil), iteration: context.iteration, configuration: configuration)
+        Event.post(.testStarted, for: (step.test, nil), configuration: configuration)
         shouldSendTestEnded = true
       case let .skip(skipInfo):
-        Event.post(.testSkipped(skipInfo), for: (step.test, nil), iteration: context.iteration, configuration: configuration)
+        Event.post(.testSkipped(skipInfo), for: (step.test, nil), configuration: configuration)
         shouldSendTestEnded = false
       case let .recordIssue(issue):
         // Scope posting the issue recorded event such that issue handling
@@ -257,7 +254,7 @@ extension Runner {
     defer {
       if let step = stepGraph.value {
         if shouldSendTestEnded {
-          Event.post(.testEnded, for: (step.test, nil), iteration: context.iteration, configuration: configuration)
+          Event.post(.testEnded, for: (step.test, nil), configuration: configuration)
         }
         Event.post(.planStepEnded(step), for: (step.test, nil), configuration: configuration)
       }
@@ -400,16 +397,37 @@ extension Runner {
   /// - Parameters:
   ///   - testCase: The test case to run.
   ///   - step: The runner plan step associated with this test case.
-  ///   - context: Context for the test run.
   ///
   /// This function sets ``Test/Case/current``, then invokes the test case's
   /// body closure.
-  private static func _runTestCase(_ testCase: Test.Case, within step: Plan.Step, context: _Context) async {
+  private static func _runTestCase(
+    _ testCase: Test.Case,
+    within step: Plan.Step,
+    context: _Context
+  ) async {
+    if _configuration.shouldUseLegacyPlanLevelRepetition {
+      await _runSingleTestCaseIteration(testCase, within: step)
+    } else {
+      await _applyRepetitionPolicy {
+        await _runSingleTestCaseIteration(testCase, within: step)
+      }
+    }
+  }
+
+  /// Run a single iteration of a test case.
+  ///
+  /// - Parameters:
+  ///   - testCase: The test case to run.
+  ///   - step: The runner plan step associated with this test case.
+  ///
+  /// This function sets ``Test/Case/current``, then invokes the test case's
+  /// body closure.
+  private static func _runSingleTestCaseIteration(_ testCase: Test.Case, within step: Plan.Step) async {
     let configuration = _configuration
 
-    Event.post(.testCaseStarted, for: (step.test, testCase), iteration: context.iteration, configuration: configuration)
+    Event.post(.testCaseStarted, for: (step.test, testCase), configuration: configuration)
     defer {
-      Event.post(.testCaseEnded, for: (step.test, testCase), iteration: context.iteration, configuration: configuration)
+      Event.post(.testCaseEnded, for: (step.test, testCase), configuration: configuration)
     }
 
     await Test.Case.withCurrent(testCase) {
@@ -434,6 +452,49 @@ extension Runner {
     }
   }
 
+  /// Applies the repetition policy specified in the current configuration by running the provided test case
+  /// repeatedly until the continuation condition is satisfied.
+  ///
+  /// - Parameters:
+  ///   - test: The test being executed.
+  ///   - testCase: The test case being iterated.
+  ///   - body: The actual body of the function which must ultimately call into the test function.
+  ///
+  /// - Note: This function updates ``Configuration/current`` before invoking the test body.
+  private static func _applyRepetitionPolicy(
+    perform body: () async -> Void
+  ) async {
+    for iteration in 1..._configuration.repetitionPolicy.maximumIterationCount {
+      let issueRecorded = Atomic(false)
+      var config = _configuration
+      config.eventHandler = { [eventHandler = config.eventHandler] event, context in
+        if case let .issueRecorded(issue) = event.kind, !issue.isKnown {
+          issueRecorded.store(true, ordering: .sequentiallyConsistent)
+        }
+        eventHandler(event, context)
+      }
+
+      await Test.withCurrentIteration(iteration) {
+        await Configuration.withCurrent(config) {
+          await body()
+        }
+      }
+
+      // Determine if the test plan should iterate again.
+      let shouldContinue = switch config.repetitionPolicy.continuationCondition {
+      case nil:
+        true
+      case .untilIssueRecorded:
+        !issueRecorded.load(ordering: .sequentiallyConsistent)
+      case .whileIssueRecorded:
+        issueRecorded.load(ordering: .sequentiallyConsistent)
+      }
+      guard shouldContinue else {
+        break
+      }
+    }
+  }
+
   /// Run the tests in this runner's plan.
   public func run() async {
     await Self._run(self)
@@ -454,21 +515,12 @@ extension Runner {
 #endif
     _ = Event.installFallbackEventHandler()
 
-    // Track whether or not any issues were recorded across the entire run.
-    let issueRecorded = Atomic(false)
-    runner.configuration.eventHandler = { [eventHandler = runner.configuration.eventHandler] event, context in
-      if case let .issueRecorded(issue) = event.kind, !issue.isKnown {
-        issueRecorded.store(true, ordering: .sequentiallyConsistent)
-      }
-      eventHandler(event, context)
-    }
-
     // Context to pass into the test run. We intentionally don't pass the Runner
     // itself (implicitly as `self` nor as an argument) because we don't want to
     // accidentally depend on e.g. the `configuration` property rather than the
     // current configuration.
     let context: _Context = {
-      var context = _Context(iteration: 0)
+      var context = _Context()
 
       let maximumParallelizationWidth = runner.configuration.maximumParallelizationWidth
       if maximumParallelizationWidth > 1 && maximumParallelizationWidth < .max {
@@ -492,45 +544,30 @@ extension Runner {
         Event.post(.runEnded, for: (nil, nil), configuration: runner.configuration)
       }
 
-      let repetitionPolicy = runner.configuration.repetitionPolicy
-      let iterationCount = repetitionPolicy.maximumIterationCount
-      for iterationIndex in 0 ..< iterationCount {
-        Event.post(.iterationStarted(iterationIndex), for: (nil, nil), configuration: runner.configuration)
-        defer {
-          Event.post(.iterationEnded(iterationIndex), for: (nil, nil), configuration: runner.configuration)
-        }
+      if runner.configuration.shouldUseLegacyPlanLevelRepetition {
+        await _applyRepetitionPolicy { [runner] in
+          let iteration = Test.currentIteration ?? 1
 
-        await withTaskGroup { [runner] taskGroup in
-          var taskAction: String?
-          if iterationCount > 1 {
-            taskAction = "running iteration #\(iterationIndex + 1)"
-          }
-          _ = taskGroup.addTaskUnlessCancelled(name: decorateTaskName("test run", withAction: taskAction)) {
-            var iterationContext = context
-            // `iteration` is one-indexed, so offset that here.
-            iterationContext.iteration = iterationIndex + 1
-            try? await _runStep(atRootOf: runner.plan.stepGraph, context: iterationContext)
+          // Legacy clients expect these values to be zero-indexed.
+          let iterationIndex = iteration - 1
+          Event.post(.iterationStarted(iterationIndex), configuration: runner.configuration)
+          defer {
+            Event.post(.iterationEnded(iterationIndex), configuration: runner.configuration)
           }
-          await taskGroup.waitForAll()
-        }
-
-        // Determine if the test plan should iterate again. (The iteration count
-        // is handled by the outer for-loop.)
-        let shouldContinue = switch repetitionPolicy.continuationCondition {
-        case nil:
-          true
-        case .untilIssueRecorded:
-          !issueRecorded.load(ordering: .sequentiallyConsistent)
-        case .whileIssueRecorded:
-          issueRecorded.load(ordering: .sequentiallyConsistent)
-        }
-        guard shouldContinue else {
-          break
+          await runner.runAll(context: context)
         }
+      } else {
+        await runner.runAll(context: context)
+      }
+    }
+  }
 
-        // Reset the run-wide "issue was recorded" flag for this iteration.
-        issueRecorded.store(false, ordering: .sequentiallyConsistent)
+  private func runAll(context: _Context) async {
+    await withTaskGroup { taskGroup in
+      _ = taskGroup.addTaskUnlessCancelled(name: decorateTaskName("test run", withAction: nil)) {
+        try? await Self._runStep(atRootOf: plan.stepGraph, context: context)
       }
+      await taskGroup.waitForAll()
     }
   }
 }

Tests/TestingTests/EventIterationTests.swift

@@ -44,7 +44,7 @@ struct EventIterationTests {
 
       // Verify all expected iterations were recorded
       let iteration = recordedIteration.load(ordering: .sequentiallyConsistent)
-      #expect(iteration == expectedIterations, "Final observed iteration did not match expected number of iterations", sourceLocation: location)
+      #expect(iteration == expectedIterations, sourceLocation: location)
     }
   }
 
@@ -60,17 +60,6 @@ struct EventIterationTests {
     }
   }
 
-  @Test
-  func `testStarted and testEnded events include iteration in context`() async {
-    await verifyIterations(
-      for: [.testStarted, .testEnded],
-      repetitionPolicy: .once,
-      expectedIterations: 1
-    ) { _ in
-      // Do nothing, just pass
-    }
-  }
-
   @Test
   func `testCaseStarted and testCaseEnded events include iteration in context`() async {
     await verifyIterations(
@@ -96,9 +85,7 @@ struct EventIterationTests {
       repetitionPolicy: policy,
       expectedIterations: expectedIterations
     ) { iteration in
-      if iteration < 3 {
-        Issue.record("Failure")
-      }
+      #expect(iteration >= 3)
     }
   }
 }

Tests/TestingTests/PlanIterationTests.swift → Tests/TestingTests/LegacyPlanIterationTests.swift

@@ -14,8 +14,8 @@
 import Synchronization
 #endif
 
-@Suite("Configuration.RepetitionPolicy Tests")
-struct PlanIterationTests {
+@Suite
+struct LegacyPlanIterationTests {
   @Test("One iteration (default behavior)")
   func oneIteration() async {
     await confirmation("N iterations started") { started in