OSGKeyboard/OSGKeyboardTests/ChunkedUtterancePipelineTests.swift

// ChunkedUtterancePipelineTests.swift
// OSGKeyboardTests

import XCTest
import os
@testable import OSGKeyboardShared

private struct StubChunkASR: ASRService, @unchecked Sendable {
    let labels: @Sendable ([Float]) -> String

    func transcribe(
        stream: AsyncStream<AudioBufferSnapshot>,
        locale: Locale
    ) -> AsyncStream<ASREvent> {
        AsyncStream { $0.finish() }
    }

    func cancel() {}

    func transcribeChunk(samples: [Float], locale: Locale) async -> ASRChunkResult {
        _ = locale
        return .success(labels(samples))
    }
}

final class ChunkedUtterancePipelineTests: XCTestCase {

    func testPipelineStitchesQueuedChunks() async {
        let config = FlowUtteranceChunkConfig(
            maxChunkDurationSeconds: 0.05,
            overlapDurationSeconds: 0,
            pauseExtensionMaxSeconds: 0,
            pauseRMSThreshold: 0.02,
            sampleRate: 1_000
        )
        let asr = StubChunkASR { samples in
            samples.isEmpty ? "" : "seg\(samples.count)"
        }
        let pipeline = ChunkedUtterancePipeline(
            asr: asr,
            locale: Locale(identifier: "zh-Hans"),
            config: config
        )

        let (stream, continuation) = AsyncStream<AudioBufferSnapshot>.makeStream()
        continuation.yield(AudioBufferSnapshot(samples: [Float](repeating: 0.1, count: 80), sampleRate: 1_000))
        continuation.yield(AudioBufferSnapshot(samples: [Float](repeating: 0.1, count: 80), sampleRate: 1_000))
        continuation.finish()

        let partialsLock = OSAllocatedUnfairLock(initialState: [String]())
        let outcome = await pipeline.transcribe(stream: stream) { partial in
            partialsLock.withLock { $0.append(partial) }
        }
        let partials = partialsLock.withLock { $0 }

        guard case .success(let success) = outcome else {
            return XCTFail("expected success, got \(outcome)")
        }
        XCTAssertTrue(success.text.contains("seg"))
        XCTAssertFalse(partials.isEmpty)
    }

    func testPipelineDeliversPartialSuccessWhenOneChunkFails() async {
        let config = FlowUtteranceChunkConfig(
            maxChunkDurationSeconds: 0.05,
            overlapDurationSeconds: 0,
            pauseExtensionMaxSeconds: 0,
            pauseRMSThreshold: 0.02,
            sampleRate: 1_000
        )
        let pipeline = ChunkedUtterancePipeline(
            asr: FailingSecondChunkASR(),
            locale: Locale(identifier: "zh-Hans"),
            config: config
        )

        let (stream, continuation) = AsyncStream<AudioBufferSnapshot>.makeStream()
        continuation.yield(AudioBufferSnapshot(samples: [Float](repeating: 0.1, count: 80), sampleRate: 1_000))
        continuation.yield(AudioBufferSnapshot(samples: [Float](repeating: 0.1, count: 80), sampleRate: 1_000))
        continuation.finish()

        let outcome = await pipeline.transcribe(stream: stream) { _ in }

        guard case .success(let success) = outcome else {
            return XCTFail("expected partial success, got \(outcome)")
        }
        XCTAssertFalse(success.text.isEmpty)
        XCTAssertEqual(success.chunkWarnings.count, 1)
    }

    func testPipelineRetranscribesShortFinalChunkWithPriorOverlap() async {
        let config = FlowUtteranceChunkConfig(
            maxChunkDurationSeconds: 0.05,
            overlapDurationSeconds: 10,
            pauseExtensionMaxSeconds: 0,
            pauseRMSThreshold: 0.02,
            minFinalChunkDurationSeconds: 0.05,
            sampleRate: 1_000
        )
        let asr = ShortFinalMergeStubASR()
        let pipeline = ChunkedUtterancePipeline(
            asr: asr,
            locale: Locale(identifier: "zh-Hans"),
            config: config
        )

        let (stream, continuation) = AsyncStream<AudioBufferSnapshot>.makeStream()
        continuation.yield(AudioBufferSnapshot(samples: [Float](repeating: 0.1, count: 80), sampleRate: 1_000))
        continuation.yield(AudioBufferSnapshot(samples: [Float](repeating: 0.1, count: 20), sampleRate: 1_000))
        continuation.finish()

        let outcome = await pipeline.transcribe(stream: stream) { _ in }
        guard case .success(let success) = outcome else {
            return XCTFail("expected success, got \(outcome)")
        }
        XCTAssertTrue(success.text.contains("merged"))
    }

    func testPipelineRetriesEmptyFinalChunkWithOverlap() async {
        let config = FlowUtteranceChunkConfig(
            maxChunkDurationSeconds: 0.05,
            overlapDurationSeconds: 10,
            pauseExtensionMaxSeconds: 0,
            pauseRMSThreshold: 0.02,
            minFinalChunkDurationSeconds: 0.05,
            sampleRate: 1_000
        )
        let asr = EmptyFinalRetryStubASR()
        let pipeline = ChunkedUtterancePipeline(
            asr: asr,
            locale: Locale(identifier: "zh-Hans"),
            config: config
        )

        let (stream, continuation) = AsyncStream<AudioBufferSnapshot>.makeStream()
        continuation.yield(AudioBufferSnapshot(samples: [Float](repeating: 0.1, count: 80), sampleRate: 1_000))
        continuation.yield(AudioBufferSnapshot(samples: [Float](repeating: 0.1, count: 80), sampleRate: 1_000))
        continuation.finish()

        let outcome = await pipeline.transcribe(stream: stream) { _ in }
        guard case .success(let success) = outcome else {
            return XCTFail("expected success, got \(outcome)")
        }
        XCTAssertTrue(success.text.contains("recovered-tail"))
    }
}

private struct FailingSecondChunkASR: ASRService, @unchecked Sendable {
    private let callIndex = OSAllocatedUnfairLock(initialState: 0)

    func transcribe(
        stream: AsyncStream<AudioBufferSnapshot>,
        locale: Locale
    ) -> AsyncStream<ASREvent> {
        AsyncStream { $0.finish() }
    }

    func cancel() {}

    func transcribeChunk(samples: [Float], locale: Locale) async -> ASRChunkResult {
        _ = locale
        let current = callIndex.withLock { state in
            let value = state
            state += 1
            return value
        }
        if current == 1 {
            return .failure("simulated chunk error")
        }
        return .success("seg\(samples.count)")
    }
}

private struct ShortFinalMergeStubASR: ASRService, @unchecked Sendable {
    private let callIndex = OSAllocatedUnfairLock(initialState: 0)

    func transcribe(
        stream: AsyncStream<AudioBufferSnapshot>,
        locale: Locale
    ) -> AsyncStream<ASREvent> {
        AsyncStream { $0.finish() }
    }

    func cancel() {}

    func transcribeChunk(samples: [Float], locale: Locale) async -> ASRChunkResult {
        _ = locale
        let current = callIndex.withLock { state in
            let value = state
            state += 1
            return value
        }
        if current == 0 {
            return .success("head")
        }
        if samples.count > 20 {
            return .success("merged-tail")
        }
        return .success("short")
    }
}

private struct EmptyFinalRetryStubASR: ASRService, @unchecked Sendable {
    private let callIndex = OSAllocatedUnfairLock(initialState: 0)

    func transcribe(
        stream: AsyncStream<AudioBufferSnapshot>,
        locale: Locale
    ) -> AsyncStream<ASREvent> {
        AsyncStream { $0.finish() }
    }

    func cancel() {}

    func transcribeChunk(samples: [Float], locale: Locale) async -> ASRChunkResult {
        _ = locale
        let current = callIndex.withLock { state in
            let value = state
            state += 1
            return value
        }
        if current == 0 {
            return .success("head")
        }
        if current == 1 {
            return .success("")
        }
        return .success("recovered-tail")
    }
}