Quick+AsyncAwait.swift

import Nimble
import Quick

/// Replacement for Quick's `it` which runs using swift concurrency.
func asyncIt(
    _ description: String,
    file: StaticString = #file,
    line: UInt = #line,
    closure: @MainActor @escaping () async throws -> Void
) {
    it(description, file: file.description, line: line) {
        innerItWait(description, file: file, line: line, closure: closure)
    }
}

/// Replacement for Quick's `fit` which runs using swift concurrency.
func fasyncIt(
    _ description: String,
    file: StaticString = #file,
    line: UInt = #line,
    closure: @MainActor @escaping () async throws -> Void
) {
    fit(description, file: file.description, line: line) {
        innerItWait(description, file: file, line: line, closure: closure)
    }
}

/// Replacement for Quick's `xit` which runs using swift concurrency.
func xasyncIt(
    _ description: String,
    file: StaticString = #file,
    line: UInt = #line,
    closure: @MainActor @escaping () async throws -> Void
) {
    xit(description, file: file.description, line: line) {
        innerItWait(description, file: file, line: line, closure: closure)
    }
}

private func innerItWait(
    _ description: String,
    file: StaticString,
    line: UInt,
    closure: @MainActor @escaping () async throws -> Void
) {
    var thrownError: Error?
    let errorHandler = { thrownError = $0 }
    let expectation = QuickSpec.current.expectation(description: description)

    Task {
        do {
            try await closure()
        } catch {
            errorHandler(error)
        }

        expectation.fulfill()
    }

    QuickSpec.current.wait(for: [expectation], timeout: 60)

    if let error = thrownError {
        XCTFail("Async error thrown: \(error)", file: file, line: line)
    }
}

/// Replacement for Quick's `beforeEach` which runs using swift concurrency.
func asyncBeforeEach(_ closure: @Sendable @MainActor @escaping (ExampleMetadata) async -> Void) {
    beforeEach({ exampleMetadata in
        let expectation = QuickSpec.current.expectation(description: "asyncBeforeEach")
        Task {
            await closure(exampleMetadata)
            expectation.fulfill()
        }
        QuickSpec.current.wait(for: [expectation], timeout: 60)
    })
}

/// Replacement for Quick's `afterEach` which runs using swift concurrency.
func asyncAfterEach(_ closure: @Sendable @MainActor @escaping (ExampleMetadata) async -> Void) {
    afterEach({ exampleMetadata in
        let expectation = QuickSpec.current.expectation(description: "asyncAfterEach")
        Task {
            await closure(exampleMetadata)
            expectation.fulfill()
        }
        QuickSpec.current.wait(for: [expectation], timeout: 60)
    })
}

/// Replacement for Nimble's `waitUntil` which waits using swift concurrency.
@discardableResult
func waitUntilAsync<R: Sendable>(
    timeout: DispatchTimeInterval = AsyncDefaults.timeout,
    action: @Sendable @escaping () async throws -> R
) async throws -> R {
    let timeInterval = timeout.timeInterval
    return try await withThrowingTaskGroup(of: R.self) { group in
        group.addTask {
            return try await action()
        }

        group.addTask {
            try await Task.sleep(nanoseconds: UInt64(timeInterval * 1_000_000_000))
            throw CancellationError()
        }

        // Starts two tasks in a group and races them until the first one finishes.
        let result = try await group.next()!
        group.cancelAll()
        return result
    }
}

/// Uses swift concurrency to poll in a loop (waiting some interval between polls) until either action returns true
/// or the timeout is reached. On timeout throws a swift `CancellationError`. Useful for waiting on one actor for
/// a result to complete in another part of the system.
func loopUntilAsync(
    timeout: DispatchTimeInterval = AsyncDefaults.timeout,
    pollInterval: DispatchTimeInterval = AsyncDefaults.pollInterval,
    action: @Sendable @escaping () async throws -> Bool
) async throws {
    return try await withThrowingTaskGroup(of: Void.self) { group in
        group.addTask {
            var hasCompleted = false
            while !hasCompleted {
                hasCompleted = try await action()
                try await Task.sleep(nanoseconds: UInt64(pollInterval.timeInterval * 1_000_000_000))
            }
        }

        group.addTask {
            try await Task.sleep(nanoseconds: UInt64(timeout.timeInterval * 1_000_000_000))
            throw CancellationError()
        }

        // Starts two tasks in a group and races them until the first one finishes.
        try await group.next()
        group.cancelAll()
    }
}

extension Expectation {

    /// Replacement for Nimble's `toEventually` which waits using swift concurrency.
    @MainActor
    func toAsyncEventually(
        _ predicate: Predicate<T>,
        timeout: DispatchTimeInterval = AsyncDefaults.timeout,
        pollInterval: DispatchTimeInterval = AsyncDefaults.pollInterval,
        description: String? = nil
    ) async {
        await innerAsyncEventually(style: .toMatch, predicate, timeout: timeout, pollInterval: pollInterval, description: description)
    }

    /// Replacement for Nimble's `toEventuallyNot` which waits using swift concurrency.
    @MainActor
    func toAsyncEventuallyNot(
        _ predicate: Predicate<T>,
        timeout: DispatchTimeInterval = AsyncDefaults.timeout,
        pollInterval: DispatchTimeInterval = AsyncDefaults.pollInterval,
        description: String? = nil
    ) async {
        await innerAsyncEventually(style: .toNotMatch, predicate, timeout: timeout, pollInterval: pollInterval, description: description)
    }

    @MainActor
    private func innerAsyncEventually(
        style: ExpectationStyle,
        _ predicate: Predicate<T>,
        timeout: DispatchTimeInterval,
        pollInterval: DispatchTimeInterval,
        description: String?
    ) async {
        let msg = FailureMessage()
        msg.userDescription = description
        msg.to = "to eventually"

        let timeInterval = pollInterval.timeInterval

        let uncachedExpression = expression.withoutCaching()

        let lastPredicateResultHolder = PredicateResultHolder()
        do {
            try await waitUntilAsync(timeout: timeout) { @MainActor in
                var hasCompleted = false
                while !hasCompleted {
                    let result = try predicate.satisfies(uncachedExpression)
                    hasCompleted = result.toBoolean(expectation: style)
                    await lastPredicateResultHolder.setLastPredicateResult(result)
                    try await Task.sleep(nanoseconds: UInt64(timeInterval * 1_000_000_000))
                }
            }
        } catch is CancellationError {
            // Async function timedout, error formatting handled as a normal completion
        } catch {
            msg.stringValue = "unexpected error thrown: <\(error)>"
            verify(false, msg)
            return
        }

        let result = (await lastPredicateResultHolder.lastPredicateResult) ?? PredicateResult(status: .fail, message: .fail("timed out before returning a value"))
        /// Note: `update` is an internal function which we've manually changed to public in our fork
        result.message.update(failureMessage: msg)
        if msg.actualValue == "" {
            msg.actualValue = "<\(stringify(try? expression.evaluate()))>"
        }
        let passed = result.toBoolean(expectation: style)
        verify(passed, msg)
    }
}

private actor PredicateResultHolder {
    var lastPredicateResult: PredicateResult?

    func setLastPredicateResult(_ result: PredicateResult) {
        lastPredicateResult = result
    }
}

fileprivate extension DispatchTimeInterval {
    var timeInterval: TimeInterval {
        switch self {
        case let .seconds(s):
            return TimeInterval(s)
        case let .milliseconds(ms):
            return TimeInterval(TimeInterval(ms) / 1000.0)
        case let .microseconds(µs):
            return TimeInterval(Int64(µs)) * TimeInterval(NSEC_PER_USEC) / TimeInterval(NSEC_PER_SEC)
        case let .nanoseconds(ns):
            return TimeInterval(ns) / TimeInterval(NSEC_PER_SEC)
        case .never:
            return .infinity
        @unknown default:
            return .infinity
        }
    }
}