swift-concurrency-extrasでローディング中(async関数のawaitの前後の状態)をTestする

動画などサイズが大きいデータをダウンロードする際に、ダウンロードのボタンが押下されたら「ダウンロード中」に、ダウンロードが完了したら「ダンロード完了」に変えるという画面があるとき、

このよう↓にダウンロードの仕組みがayncで提供されているViewModelなどで、awaitの前後でtextが「ダウンロード中」→「ダウンロード完了」に適切に切り替わることをテストしたいです
このテストの方法として、自分がやっていたパワー？なCheckedContinuationを管理する方法とswift-concurrency-extrasを使った方法を紹介します

class ContentViewModel: ObservableObject {
    @Published var text: String = "ダウンロードする"

    private let downloadHandler: () async -> String
    init(downloadHandler: @escaping () async -> String) {
        self.downloadHandler = downloadHandler
    }

    func download() async {
        text = "ダウンロード中" // ← 関数が呼ばれてからawait中の状態
        let res = await downloadHandler()
        text = "ダウンロード完了: \(res)" // ← await後の状態
    }
}

(Xcode14.3で検証しています)

CheckedContinuationを管理する方法

最初はこんな感じ↓のテストを実装しました
しかし、この実装ではtextが「ダウンロード中」を期待していたタイミングで、「ダウンロードする」から変更されておらずエラーとなってしまいました

class ContentViewModelTests: XCTestCase {
    func testDownload() async {
        let vm = ContentViewModel { "Hello" }
        let task = Task { await vm.download() }

        // vm.textが"ダウンロードする"のままで、テストが失敗する
        XCTAssertEqual(vm.text, "ダウンロード中") // XCTAssertEqual failed: ("ダウンロードする") is not equal to ("ダウンロード中")

        await task.value

        XCTAssertEqual(vm.text, "ダウンロード完了: Hello")
    }
}

これは、Taskの中身の未実行が問題なので、Task.sleepを挟むことで解決できました
しかし、これではテストにかかる時間が増えてしまいます

class ContentViewModelTests: XCTestCase {
    func testDownload() async {
        let vm = ContentViewModel {
+           try? await Task.sleep(for: .seconds(1))
            return "Hello"
        }
        let task = Task { await vm.download() }

+       try? await Task.sleep(for: .seconds(0.5))

        XCTAssertEqual(vm.text, "ダウンロード中")

        await task.value

        // テストの完了まで時間がかかってしまう
        XCTAssertEqual(vm.text, "ダウンロード完了: Hello")
    }
}

そこで、CurrentValueSubjectを使い、withCheckedContinuationでCheckedContinuationが生成されるまで待機させることで、downloadHandlerが実行されるまで待機するのを実現しました
この方法ではテストは問題なく通過しますし、任意のタイミングでイベントを流せて便利なのですが、ボイラープレートが多いことやCheckedContinuationに1回だけイベントを流すのを実装で保証しないといけないことなどの課題感を持っていました

class ContentViewModelTests: XCTestCase {
    func testDownload() async {
+       let subject = CurrentValueSubject<CheckedContinuation<String, Never>?, Never>(nil)

        let vm = ContentViewModel {
+           await withCheckedContinuation { (c: CheckedContinuation<String, Never>) in
+               // downloadHandlerが実行されたらsubjectに通知
+               subject.send(c)
+           }
-           "Hello"
        }
        let task = Task { await vm.download() }

+       // downloadHandlerが実行されるまで待機
+       let continuation = await subject.values.first(where: { $0 != nil })!!

        XCTAssertEqual(vm.text, "ダウンロード中")

+       continuation.resume(with: .success("Hello"))
        await task.value

        XCTAssertEqual(vm.text, "ダウンロード完了: Hello")
    }
}

swift-concurrency-extrasを使う方法

このローディング中のテストは、pointfreecoが公開しているswift-concurrency-extrasに入っているwithMainSerialExecutorを使ってシンプルに実現できました

テストをwithMainSerialExecutorに入れ、Task生成後とdownloadHandlerで値を返す前にTask.yield()を入れると、テストが通過するようになります
Task.yield()は増えましたが、CheckedContinuationの管理がなくなり、シンプルにテストを実現できたと思います

class ContentViewModelTests: XCTestCase {
    func testDownload() async {
+       await withMainSerialExecutor {
            let vm = ContentViewModel {
+               await Task.yield()
                return "Hello"
            }
            let task = Task { await vm.download() }

+           await Task.yield()

            XCTAssertEqual(vm.text, "ダウンロード中")

            await task.value

            XCTAssertEqual(vm.text, "ダウンロード完了: Hello")
+       }
    }
}

もし、withMainSerialExecutorを外し、ContentViewModelに@.MainActorをつけると、一瞬テストが通過するような感じになるのですが、1000回くらい繰り返すと確率で失敗するFlakyなテストとなってしまいました
withMainSerialExecutorはちゃんと何か重要な役割を果たしていそうです

class ContentViewModelTests: XCTestCase {
+   @MainActor
    func testDownload() async {
-       await withMainSerialExecutor {
            let vm = ContentViewModel {
                await Task.yield()
                return "Hello"
            }
            let task = Task { await vm.download() }

            await Task.yield()

            // 確率で失敗してしまう
            XCTAssertEqual(vm.text, "ダウンロード中")

            await task.value

            XCTAssertEqual(vm.text, "ダウンロード完了: Hello")
-       }
    }
}

withMainSerialExecutorの実装を見てみる

withMainSerialExecutorの実装を少しだけ見てみました
大事そうな部分の抜粋です
swift_task_enqueueGlobal_hookというのを取得し、MainActorにenqueueさせるようにしているようです
これにより、全てのTaskがMainActorで実行されることで、順序の問題がなくなり、Flakyな問題が解消されるようです(多分)

@MainActor
public func withMainSerialExecutor(
    @_implicitSelfCapture operation: @MainActor @Sendable () async throws -> Void
) async rethrows {
    let didUseMainSerialExecutor = uncheckedUseMainSerialExecutor
    defer { uncheckedUseMainSerialExecutor = didUseMainSerialExecutor }
    uncheckedUseMainSerialExecutor = true
    try await operation()
}
// 中略
public var uncheckedUseMainSerialExecutor: Bool {
    get { swift_task_enqueueGlobal_hook != nil }
    set {
        swift_task_enqueueGlobal_hook =
            newValue ? { job, _ in MainActor.shared.enqueue(job) } : nil
    }
}
// 中略
private let _swift_task_enqueueGlobal_hook: UnsafeMutablePointer<Hook?> =
  dlsym(dlopen(nil, 0), "swift_task_enqueueGlobal_hook").assumingMemoryBound(to: Hook?.self)

まとめ

Swift Concurrencyでローディング中をTestするために、自分がやっていたCheckedContinuationを管理する方法とswift-concurrency-extrasを使った方法を紹介しました
swift-concurrency-extrasのwithMainSerialExecutorを使うことで、シンプルにテストができました
ただ、CheckedContinuationを管理する方法では任意のタイミングでイベントを流せるというメリットもあるので、いい感じに組み合わせてテストが実装できると良さそうです
(公式でこの辺をいい感じにやる仕組みを作って欲しいという気持ちもあったりします)