Async/Await单元测试失败：ViewModel中Do代码块未执行求助

我正在编写基于async/await的单元测试，尝试多种组合后仍无法通过测试。调用ViewModel的checkNumber()方法时，仅执行前两行代码，Do代码块被跳过，但集成测试无此问题，希望排查原因。

测试代码 #

func test_checkNumber_success() async {
    let expectation = XCTestExpectation(description: "CheckNumber completes successfully")
    
    viewModel = PrimeNumberCheckViewModel(service: MockPrimeNumberService())
    
    XCTAssertEqual(viewModel.isPrime, false)
    XCTAssertEqual(viewModel.validationInProgress, false)
    
    let res = await viewModel.checkNumber(3)
    
    expectation.fulfill()
    
    XCTAssertEqual(viewModel.isPrime, true)
    XCTAssertEqual(viewModel.validationInProgress, false)
    
    wait(for: [expectation], timeout: 0.05) // Adjust the timeout as needed
}

ViewModel代码 #

final class PrimeNumberCheckViewModel: PrimeNumberCheckProtocol {
    @Published private(set) var validationInProgress =  false
    @Published var status = "checking ready to get started"
    @Published private(set)var isPrime: Bool = false
    
    private let service: PrimeServiceProtocol
    private var validationTask: Task<Void, Never>?
    
    init(service: PrimeServiceProtocol) {
        self.service = service
    }
    
    @MainActor
    func checkNumber(_ number: Int) async {

        self.validationInProgress = true
        self.status = "Checking \(number) in progress"
        
        validationTask?.cancel()
        
        validationTask = Task {
            var res = false

            do {
                res = try await service.checkNumber(number)
                print("this is res: \(res)")
                
                updateStatus(res, number)
                
                self.isPrime = res
            } catch {
                //TODO: handle error scenario case
                print(error)
            }
            self.validationInProgress = false
        }
    }
}

Mock服务代码 #

class MockPrimeNumberService: PrimeServiceProtocol {
    @MainActor
    func checkNumber(_ number: Int) async throws -> Bool {
        try await Task.sleep(nanoseconds: 2 * 1_000_000_000)
        
        // Check if the number is less than 2
        if number <= 1 {
            return false
        }
        
        // Check for divisibility from 2 to the square root of the number
        for i in 2..<Int(Double(number).squareRoot()) + 1 {
            if number % i == 0 {
                return false
            }
        }
        return true
    }
}

问题根源 #

1. 测试断言时机过早：checkNumber方法内部启动了独立Task后会立即返回，测试代码在调用await viewModel.checkNumber(3)后立刻执行断言，此时内部Task还未开始执行Do代码块，导致断言失败。
2. Expectation使用错误：测试中调用checkNumber后立刻调用expectation.fulfill()，没有等待内部Task完成；且设置的0.05秒超时远小于Mock服务的2秒休眠，直接导致测试提前结束。
3. Mock服务冗余延迟：单元测试不需要真实的时间延迟，2秒休眠不仅拖慢测试，还会导致测试超时。

修复方案 #

1. 调整测试逻辑，正确等待Task完成

保留原ViewModel设计的前提下，通过监听Published值的变化来等待任务结束：

func test_checkNumber_success() async {
    viewModel = PrimeNumberCheckViewModel(service: MockPrimeNumberService())
    
    XCTAssertEqual(viewModel.isPrime, false)
    XCTAssertEqual(viewModel.validationInProgress, false)
    
    await viewModel.checkNumber(3)
    
    // 监听validationInProgress状态变为false，确认任务完成
    let completionExpectation = XCTestExpectation(description: "Validation finishes")
    let cancellable = viewModel.$validationInProgress.sink { isInProgress in
        if !isInProgress {
            completionExpectation.fulfill()
        }
    }
    
    wait(for: [completionExpectation], timeout: 2.1)
    
    XCTAssertEqual(viewModel.isPrime, true)
    XCTAssertEqual(viewModel.validationInProgress, false)
    
    cancellable.cancel()
}

2. 移除Mock服务的冗余延迟（推荐）

单元测试应快速执行，删除Task.sleep：

class MockPrimeNumberService: PrimeServiceProtocol {
    @MainActor
    func checkNumber(_ number: Int) async throws -> Bool {
        if number <= 1 {
            return false
        }
        
        for i in 2..<Int(Double(number).squareRoot()) + 1 {
            if number % i == 0 {
                return false
            }
        }
        return true
    }
}

3. 可选：修改ViewModel让方法等待任务完成

如果业务允许，可以让checkNumber直接执行逻辑而非启动独立Task，避免异步等待问题：

@MainActor
func checkNumber(_ number: Int) async {
    validationTask?.cancel()
    
    self.validationInProgress = true
    self.status = "Checking \(number) in progress"
    
    var res = false
    do {
        res = try await service.checkNumber(number)
        print("this is res: \(res)")
        
        updateStatus(res, number)
        
        self.isPrime = res
    } catch {
        print(error)
    }
    self.validationInProgress = false
}

内容的提问来源于stack exchange，提问作者dioufism