OSGKeyboard/OSGKeyboardTests/LocalASRModelCatalogTests.swift

// LocalASRModelCatalogTests.swift
// OSGKeyboardTests

import XCTest
@testable import OSGKeyboardShared

final class LocalASRModelCatalogTests: XCTestCase {

    func testBundledCatalogLoads() throws {
        let catalog = try LocalASRModelCatalog.loadBundled()
        XCTAssertEqual(catalog.schemaVersion, 1)
        XCTAssertEqual(catalog.defaultModelId, "sherpa-qwen3-0.6b-int8")
        XCTAssertFalse(catalog.models.contains { $0.id == "qwen3-mlx-1.7b" })
        XCTAssertTrue(catalog.models.contains { $0.id == "sherpa-qwen3-0.6b-int8" })
        XCTAssertTrue(catalog.models.contains { $0.id == "sherpa-qwen3-1.7b-int8" })
        XCTAssertTrue(catalog.models.contains { $0.id == "sherpa-sensevoice-small-int8" })
        XCTAssertFalse(catalog.models.contains { $0.id == "sherpa-paraformer-zh-int8" })
        XCTAssertEqual(
            LocalASRModelCatalog.model("sherpa-sensevoice-small-int8", in: catalog)?.badgeKey,
            "mac.localASR.badge.fastest"
        )
        XCTAssertEqual(
            LocalASRModelCatalog.model("sherpa-qwen3-0.6b-int8", in: catalog)?.badgeKey,
            "mac.localASR.badge.balanced"
        )
        XCTAssertEqual(
            LocalASRModelCatalog.model("sherpa-qwen3-1.7b-int8", in: catalog)?.badgeKey,
            "mac.localASR.badge.quality"
        )
    }

    func testSherpaQwen317BUsesRepositoryInstall() throws {
        let catalog = try LocalASRModelCatalog.loadBundled()
        let model = try XCTUnwrap(LocalASRModelCatalog.model("sherpa-qwen3-1.7b-int8", in: catalog))
        XCTAssertEqual(model.installKind, .repository)
        XCTAssertTrue(model.sources?.contains(where: { $0.type == "modelscope" && $0.isRepository }) == true)
        XCTAssertTrue(model.sources?.contains(where: { $0.type == "huggingface" && $0.isRepository }) == true)
    }

    func testCapabilitiesForSherpaQwen3() throws {
        let catalog = try LocalASRModelCatalog.loadBundled()
        let model = try XCTUnwrap(LocalASRModelCatalog.model("sherpa-qwen3-0.6b-int8", in: catalog))
        let caps = LocalASRModelCatalog.capabilities(for: model)
        XCTAssertEqual(caps.hotwordMode, .recognizerScoped)
        XCTAssertTrue(model.supportsHotwords)
    }

    func testCapabilitiesForSenseVoice() throws {
        let catalog = try LocalASRModelCatalog.loadBundled()
        let model = try XCTUnwrap(LocalASRModelCatalog.model("sherpa-sensevoice-small-int8", in: catalog))
        let caps = LocalASRModelCatalog.capabilities(for: model)
        XCTAssertEqual(caps.hotwordMode, .none)
        XCTAssertFalse(model.supportsHotwords)
    }

    func testManifestRoundTrip() throws {
        let manifest = LocalASRInstalledManifest(
            selectedModelId: "sherpa-qwen3-0.6b-int8",
            installedModelIDs: ["sherpa-qwen3-0.6b-int8"]
        )
        let url = FileManager.default.temporaryDirectory
            .appendingPathComponent("manifest-\(UUID().uuidString).json")
        defer { try? FileManager.default.removeItem(at: url) }

        let encoder = JSONEncoder()
        encoder.dateEncodingStrategy = .iso8601
        try encoder.encode(manifest).write(to: url)

        let decoder = JSONDecoder()
        decoder.dateDecodingStrategy = .iso8601
        let loaded = try decoder.decode(LocalASRInstalledManifest.self, from: Data(contentsOf: url))
        XCTAssertEqual(loaded.selectedModelId, manifest.selectedModelId)
        XCTAssertEqual(loaded.installedModelIDs, manifest.installedModelIDs)
    }

    func testBiasDiagnosticsStoreRoundTrip() {
        LocalASRBiasDiagnosticsStore.clear()
        let payload = LocalASRBiasPayload(
            hardHotwords: ["Cursor"],
            promptBias: "test",
            corpusContext: nil,
            polishFragment: "fragment",
            correctionPairs: [],
            diagnostics: LocalASRBiasDiagnostics(userTermCount: 2, builtinTermCount: 3)
        )
        LocalASRBiasDiagnosticsStore.save(
            payload: payload,
            modelId: "sherpa-qwen3-0.6b-int8",
            backendLabel: "Sherpa Qwen3"
        )
        let snapshot = LocalASRBiasDiagnosticsStore.load()
        XCTAssertEqual(snapshot?.modelId, "sherpa-qwen3-0.6b-int8")
        XCTAssertEqual(snapshot?.diagnostics.userTermCount, 2)
        XCTAssertEqual(snapshot?.hotwordCount, 1)
        LocalASRBiasDiagnosticsStore.clear()
    }

    func testSherpaAdapterProducesHardHotwords() throws {
        let fixtureURL = FileManager.default.temporaryDirectory
            .appendingPathComponent("phrases-\(UUID().uuidString).tsv")
        try "word\tpinyin\tsource\tweight\nSwiftUI\tswift ui\tcomputer_terms\t5\n"
            .write(to: fixtureURL, atomically: true, encoding: .utf8)
        defer { try? FileManager.default.removeItem(at: fixtureURL) }

        var dict = PersonalDictionary.empty
        _ = dict.upsertManual(term: "Kubernetes")

        let payload = LocalASRBiasAdapter.adapt(
            LocalASRBiasRequest(
                dictionary: dict,
                locale: Locale(identifier: "zh-CN"),
                capabilities: .sherpaQwen3
            ),
            lexicon: BuiltinLexiconIndex(fixtureURL: fixtureURL)
        )
        XCTAssertFalse(payload.hardHotwords.isEmpty)
        XCTAssertTrue(payload.hardHotwords.contains("Kubernetes"))
    }
}