Swift Testingでconfirmationがconfirmされるまで待つ

前提

以下のような非同期関数（Swift Concurency以前の物や、どうしてもcompletion形式になってしまう物）をテストしたいとする。

func doSomething(completion: @Sendable @escaping () -> Void) {
    Task {
        try await Task.sleep(for: .seconds(0.1)) // なんか重い処理
        completion()
    }
}

このような場合、XCTestではexpectationを使っていた。

func testDoSomething() async {
    let exp = expectation(description: "doSomething")
    doSomething {
        exp.fulfill()
    }
    
    await fulfillment(of: [exp], timeout: 10.0)
}

Swift Testingではconfirmationを使う。

@Test
func testDoSomething() async {
    await confirmation(expectedCount: 1) { confirm in
        doSomething {
            confirm()
        }
    }
}

しかしこのテストは失敗する。なぜならconfirmationはconfirm()が実行されるまでsuspendしてくれないからである。（await fulfillment(of: [exp], timeout: 10.0)に該当する構文がSwift Testingにない。）

https://developer.apple.com/documentation/testing/testing-asynchronous-code

Confirmations function similarly to the expectations API of XCTest, however, they don’t block or suspend the caller while waiting for a condition to be fulfilled. Instead, the requirement is expected to be confirmed (the equivalent of fulfilling an expectation) before confirmation() returns, and records an issue otherwise:
Confirmationは XCTest のExpectation API と同様に機能するが、条件が満たされるのを待つ間、呼び出し元をブロックしたり中断したりしない。その代わり、confirmation() が終了する前にconfirm()が実行される（fullfill()と同じ）ことが期待され、そうでない場合は問題が記録される。

解決策

元も子もないが、最も簡単な解決策はcompletion-styleの非同期関数をasync化することである。
しかし、それができない場合は以下のような実装を使った。

[追記] バグがあったためicemanさんのコメントを元にコードを微修正

func suspendUntilConfirmation(
    expectedCount: Int = 1,
    timeout: Duration? = nil,
    isolation: isolated (any Actor)? = #isolation,
    body: (@escaping @Sendable () -> Void) -> Void
) async {
    await confirmation(expectedCount: expectedCount) { confirm in
        var timeoutTask: Task<Void, Error>?
        defer {
            timeoutTask?.cancel()
        }

        // confirmされるまでsuspendする
        await withCheckedContinuation { (continuation: CheckedContinuation<Void, Never>) in
            let continuation: OSAllocatedUnfairLock<CheckedContinuation<Void, Never>?> = .init(
                initialState: continuation
            )
            // Confirmation.countがprivateなので自分でカウンタを持つ
            let count: OSAllocatedUnfairLock<Int> = .init(initialState: 0)
            body {
                let countResult = count.withLock {
                    $0 += 1
                    return $0
                }
                confirm()

                if countResult == expectedCount {
                    continuation.withLock {
                        $0?.resume()
                        $0 = nil
                    }
                }
            }

            if let timeout {
                timeoutTask = Task {
                    try await Task.sleep(for: timeout)
                    continuation.withLock {
                        $0?.resume()
                        $0 = nil
                    }
                }
            }
        }
    }
}

func testDoSomething() async {
    await suspendUntilConfirmation(expectedCount: 1) { confirm in
        doSomething {
            confirm()
        }
    }

    // タイムアウト指定もできる！
    await suspendUntilConfirmation(expectedCount: 1, timeout: .seconds(0.3)) { confirm in
        doSomething {
            confirm()
        }
    }
}

補足

closureがSendableになっても良いならTaskGroupの実装の方が綺麗

func suspendUntilConfirmation(
    expectedCount: Int = 1,
    timeout: Duration? = nil,
    body: @escaping @Sendable (@escaping @Sendable () -> Void) -> Void
) async {
    await confirmation(expectedCount: expectedCount) { confirm in
        await withTaskGroup(of: Void.self) { group in
            let continuation: OSAllocatedUnfairLock<CheckedContinuation<Void, Never>?> = .init(initialState: nil)
            let count: OSAllocatedUnfairLock<Int> = .init(initialState: 0)

            group.addTask {
                await withCheckedContinuation { cont in
                    continuation.withLock { $0 = cont }

                    body {
                        let countResult = count.withLock {
                            $0 += 1
                            return $0
                        }
                        confirm()
            
                        if countResult == expectedCount {
                            continuation.withLock {
                                $0?.resume()
                                $0 = nil
                            }
                        }
                    }
                }
            }

            if let timeout {
                group.addTask {
                    try? await Task.sleep(for: timeout)
                    continuation.withLock {
                        $0?.resume()
                        $0 = nil
                    }
                }
            }
            
            await group.next()
            group.cancelAll()
        }
    }
}

余談

「confirmされるまでsuspendする」ができないというのは必要なパーツが欠けている気がしていて、ただ自分が勘違いしている/使い方を間違えてるだけかもしれないので、良いアプローチを知っている人がいたらコメントください。

せめてConfirmation.countが触れるようにするとか、await fullfill(confirm:count:)みたいな関数を追加するとかしてほしいが

await confirmation { confirm in
    body {
        confirm()
    }

    // とか
    while confirm.count != 1 {
        await Task.yield()
    }

    // とか
    await fullfill(confirm, count: 1)
}

というかドキュメントに載っているconfirmationの使用例コード、元のcompletionの内部実装によってはflakyなテストな気がするんだが...?

参考

タイムアウトの実装

https://zenn.dev/iceman/articles/8e863b48b035b3#タイムアウトを設ける