XCTestCase+Wait.swift

import Foundation
import XCTest

extension XCTestCase {
    /// Waits for an escaping closure to be invoked with a result value for the specified timeout duration and returns the value.
    ///
    /// When using this method, you must trigger work inside the `performWork` closure that invokes the `completion` closure exactly once.
    /// Failure to do so within the specified `timeout` will result in an error being thrown.
    ///
    /// ```swift
    /// func testSomething() throws {
    ///     // Given a service
    ///     let service = makeService()
    ///
    ///     // When the service performs an asynchronous request
    ///     let result = try wait(timeout: 0.1) { fulfill in
    ///         service.fetchValueAsynchronously { fulfill($0) }
    ///     }
    ///
    ///     // Then the result will match the expected outcome
    ///     XCTAssertEqual(result, .success(.fake(id: 1))) // `result` is the value returned by `fetchValueAsynchronously(_:)`'s closure
    /// }
    /// ```
    ///
    /// - Parameters:
    ///   - timeout: The timeout duration (in seconds) of the test expectation that waits for the result.
    ///   - description: A description used for the underlying text expectation. Defaults to the caller `#function` name.
    ///   - performWork: A non-escaping closure used to trigger the asynchronous work that is being awaited.
    ///     Provides one parameter which is a fulfilment closure to be called once the underlying work has finished and a result has been returned.
    /// - Returns: The value passed into the completion closure of the `performWork` closure.
    /// - Throws: If `performWork`'s completion closure was not invoked before the specified timeout.
    public func wait<T>(
        timeout: TimeInterval,
        description: String = #function,
        for performWork: (@escaping (T) -> Void) -> Void
    ) throws -> T {
        let expectation = expectation(description: description)
        var result: T?

        performWork {
            result = $0
            expectation.fulfill()
        }

        // Actually wait for the expectation to finish and attempt to return the result
        wait(for: [expectation], timeout: timeout)
        return try XCTUnwrap(result, "completion was not called within the \(timeout)s timeout duration")
    }

    /// Waits for an escaping closure to be invoked for the specified timeout duration.
    ///
    /// - Parameters:
    ///   - timeout: The timeout duration (in seconds) of the test expectation that waits for the result.
    ///   - description: A description used for the underlying text expectation. Defaults to the caller `#function` name.
    ///   - performWork: A non-escaping closure used to trigger the asynchronous work that is being awaited.
    ///     Provides one parameter which is a fulfilment closure to be called once the underlying work has finished.
    /// - Throws: If `performWork`'s completion closure was not invoked before the specified timeout.
    public func wait(
        timeout: TimeInterval,
        description: String = #function,
        for performWork: (@escaping () -> Void) -> Void
    ) throws {
        try wait(timeout: timeout, description: description) { fulfill in
            performWork { fulfill(()) }
        }
    }

    /// Waits for a queue to fulfill an expectation asynchronously or times out.
    ///
    /// When the provided `queue` is serial, any async work previously scheduled will have to complete allowing the expectation to wait for that work to complete before fulfilling.
    ///
    /// Use this method with serial queues such as `DispatchQueue.main` in order to allow other pending work items to be performed to essentially flush a queue before making further assertions. For example:
    ///
    /// ```swift
    /// var isFinished: Bool = false
    /// DispatchQueue.main.async {
    ///     isFinished = true
    /// }
    ///
    /// XCTAssertFalse(isFinished) // ✅
    /// wait(for: .main, timeout: 0.1)
    /// XCTAssertTrue(isFinished) // ✅
    /// ```
    ///
    /// - Parameters:
    ///   - queue: The queue wait on
    ///   - timeout: The timeout duration in seconds to wait until a failure is triggered
    public func wait(for queue: DispatchQueue, timeout: TimeInterval) {
        // Create an expectation that will be used to wait for the queue to perform an async closure
        let expectation = self.expectation(description: "DispatchQueue async expectation")

        // Schedule the expectation on the queue
        queue.async {
            expectation.fulfill()
        }

        // Wait for the `XCTestExpectation` to fulfill or timeout
        wait(for: [expectation], timeout: timeout)
    }
}