Implement XCTestCase.addTeardownBlock

Sources/XCTest/Public/XCTestCase.swift

@@ -132,6 +132,7 @@ open class XCTestCase: XCTest {
                     expected: false)
             }
         }
+        runTeardownBlocks()
         tearDown()
     }
 
@@ -178,6 +179,32 @@ open class XCTestCase: XCTest {
     /// class.
     open class func tearDown() {}
 
+    private var teardownBlocks: [() -> Void] = []
+    private var teardownBlocksDequeued: Bool = false
+    private let teardownBlocksQueue: DispatchQueue = DispatchQueue(label: "org.swift.XCTest.XCTestCase.teardownBlocks")
+
+    /// Registers a block of teardown code to be run after the current test
+    /// method ends.
+    open func addTeardownBlock(_ block: @escaping () -> Void) {
+        teardownBlocksQueue.sync {
+            precondition(!self.teardownBlocksDequeued, "API violation -- attempting to add a teardown block after teardown blocks have been dequeued")
+            self.teardownBlocks.append(block)
+        }
+    }
+
+    private func runTeardownBlocks() {
+        let blocks = teardownBlocksQueue.sync { () -> [() -> Void] in
+            self.teardownBlocksDequeued = true
+            let blocks = self.teardownBlocks
+            self.teardownBlocks = []
+            return blocks
+        }
+
+        for block in blocks.reversed() {
+            block()
+        }
+    }
+
     open var continueAfterFailure: Bool {
         get {
             return true

Tests/Functional/TestCaseLifecycle/Misuse/main.swift

@@ -0,0 +1,40 @@
+// RUN: %{swiftc} %s -g -o %T/TestCaseLifecycleMisuse
+// RUN: %T/TestCaseLifecycleMisuse > %t || true
+// RUN: %{xctest_checker} %t %s
+
+#if os(macOS)
+    import SwiftXCTest
+#else
+    import XCTest
+#endif
+
+// CHECK: Test Suite 'All tests' started at \d+-\d+-\d+ \d+:\d+:\d+\.\d+
+// CHECK: Test Suite '.*\.xctest' started at \d+-\d+-\d+ \d+:\d+:\d+\.\d+
+
+// CHECK: Test Suite 'TeardownBlocksMisuseTestCase' started at \d+-\d+-\d+ \d+:\d+:\d+\.\d+
+class TeardownBlocksMisuseTestCase: XCTestCase {
+
+    // CHECK: Test Case 'TeardownBlocksMisuseTestCase.test_addingATeardownBlockLate_crashes' started at \d+-\d+-\d+ \d+:\d+:\d+\.\d+
+    func test_addingATeardownBlockLate_crashes() {
+        addTeardownBlock { [weak self] in
+            guard let self = self else {
+                XCTFail("self should still exist at this point")
+                return
+            }
+            // The following line should crash and nothing more will be printed
+            self.addTeardownBlock {
+                print("This should not be printed")
+            }
+        }
+    }
+
+    static var allTests = {
+      return [
+          ("test_addingATeardownBlockLate_crashes", test_addingATeardownBlockLate_crashes),
+      ]
+    }()
+}
+
+XCTMain([
+    testCase(TeardownBlocksMisuseTestCase.allTests),
+])

Tests/Functional/TestCaseLifecycle/main.swift

@@ -86,12 +86,62 @@ class NewInstanceForEachTestTestCase: XCTestCase {
 // CHECK: \t Executed 2 tests, with 0 failures \(0 unexpected\) in \d+\.\d+ \(\d+\.\d+\) seconds
 
 
+// CHECK: Test Suite 'TeardownBlocksTestCase' started at \d+-\d+-\d+ \d+:\d+:\d+\.\d+
+class TeardownBlocksTestCase: XCTestCase {
+    static var allTests = {
+        return [
+            ("test_withoutTeardownBlocks", test_withoutTeardownBlocks),
+            ("test_withATeardownBlock", test_withATeardownBlock),
+            ("test_withSeveralTeardownBlocks", test_withSeveralTeardownBlocks),
+        ]
+    }()
+
+    override func tearDown() {
+        print("In tearDown function")
+    }
+
+    // CHECK: Test Case 'TeardownBlocksTestCase.test_withoutTeardownBlocks' started at \d+-\d+-\d+ \d+:\d+:\d+\.\d+
+    // CHECK: In tearDown function
+    // CHECK: Test Case 'TeardownBlocksTestCase.test_withoutTeardownBlocks' passed \(\d+\.\d+ seconds\)
+    func test_withoutTeardownBlocks() {
+        // Intentionally blank
+    }
+
+    // CHECK: Test Case 'TeardownBlocksTestCase.test_withATeardownBlock' started at \d+-\d+-\d+ \d+:\d+:\d+\.\d+
+    // CHECK: In teardown block A
+    // CHECK: In tearDown function
+    // CHECK: Test Case 'TeardownBlocksTestCase.test_withATeardownBlock' passed \(\d+\.\d+ seconds\)
+    func test_withATeardownBlock() {
+        addTeardownBlock {
+            print("In teardown block A")
+        }
+    }
+
+    // CHECK: Test Case 'TeardownBlocksTestCase.test_withSeveralTeardownBlocks' started at \d+-\d+-\d+ \d+:\d+:\d+\.\d+
+    // CHECK: In teardown block C
+    // CHECK: In teardown block B
+    // CHECK: In tearDown function
+    // CHECK: Test Case 'TeardownBlocksTestCase.test_withSeveralTeardownBlocks' passed \(\d+\.\d+ seconds\)
+    func test_withSeveralTeardownBlocks() {
+        addTeardownBlock {
+            print("In teardown block B")
+        }
+        addTeardownBlock {
+            print("In teardown block C")
+        }
+    }
+}
+// CHECK: Test Suite 'TeardownBlocksTestCase' passed at \d+-\d+-\d+ \d+:\d+:\d+\.\d+
+// CHECK: \t Executed 3 tests, with 0 failures \(0 unexpected\) in \d+\.\d+ \(\d+\.\d+\) seconds
+
+
 XCTMain([
     testCase(SetUpTearDownTestCase.allTests),
-    testCase(NewInstanceForEachTestTestCase.allTests)
+    testCase(NewInstanceForEachTestTestCase.allTests),
+    testCase(TeardownBlocksTestCase.allTests),
 ])
 
 // CHECK: Test Suite '.*\.xctest' passed at \d+-\d+-\d+ \d+:\d+:\d+\.\d+
-// CHECK: \t Executed 3 tests, with 0 failures \(0 unexpected\) in \d+\.\d+ \(\d+\.\d+\) seconds
+// CHECK: \t Executed 6 tests, with 0 failures \(0 unexpected\) in \d+\.\d+ \(\d+\.\d+\) seconds
 // CHECK: Test Suite 'All tests' passed at \d+-\d+-\d+ \d+:\d+:\d+\.\d+
-// CHECK: \t Executed 3 tests, with 0 failures \(0 unexpected\) in \d+\.\d+ \(\d+\.\d+\) seconds
+// CHECK: \t Executed 6 tests, with 0 failures \(0 unexpected\) in \d+\.\d+ \(\d+\.\d+\) seconds