test_benchmarking.py

import unittest

from benchmarking import (
    BENCHMARK_MODEL_SPECS,
    DEFAULT_BENCHMARK_MODEL_IDS,
    build_benchmark_row,
    get_benchmark_specs,
    prepare_benchmark_text,
)


class BenchmarkingTest(unittest.TestCase):
    def test_default_models_include_current_adapter_and_fast_baselines(self):
        self.assertIn("current-adapter", DEFAULT_BENCHMARK_MODEL_IDS)
        self.assertIn("multilingual-e5-small", DEFAULT_BENCHMARK_MODEL_IDS)
        self.assertIn("multilingual-minilm", DEFAULT_BENCHMARK_MODEL_IDS)

    def test_korean_focused_models_include_kosimcse_and_kure(self):
        specs_by_key = {spec.key: spec for spec in BENCHMARK_MODEL_SPECS}

        self.assertEqual(
            specs_by_key["kosimcse-roberta-multitask"].model_id,
            "BM-K/KoSimCSE-roberta-multitask",
        )
        self.assertEqual(
            specs_by_key["kosimcse-roberta-multitask"].kind,
            "transformers_cls",
        )
        self.assertEqual(
            specs_by_key["kr-sbert"].model_id,
            "snunlp/KR-SBERT-V40K-klueNLI-augSTS",
        )
        self.assertEqual(specs_by_key["kure-v1"].model_id, "nlpai-lab/KURE-v1")

    def test_get_benchmark_specs_preserves_catalog_order_and_ignores_unknown(self):
        selected = ["unknown", "multilingual-minilm", "current-adapter"]

        specs = get_benchmark_specs(selected)

        self.assertEqual(
            [spec.model_id for spec in specs],
            [
                BENCHMARK_MODEL_SPECS[0].model_id,
                "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2",
            ],
        )

    def test_prepare_benchmark_text_adds_e5_prefix_only_for_e5_models(self):
        e5_spec = get_benchmark_specs(["multilingual-e5-small"])[0]
        current_spec = get_benchmark_specs(["current-adapter"])[0]

        self.assertEqual(
            prepare_benchmark_text(e5_spec, "가맹점명: 테스트"),
            "query: 가맹점명: 테스트",
        )
        self.assertEqual(
            prepare_benchmark_text(current_spec, "가맹점명: 테스트"),
            "가맹점명: 테스트",
        )

    def test_build_benchmark_row_formats_score_elapsed_and_status(self):
        spec = get_benchmark_specs(["current-adapter"])[0]

        row = build_benchmark_row(
            spec,
            score=0.912345,
            judgement="강유사 후보",
            elapsed_ms=123.456,
            status="OK",
        )

        self.assertEqual(
            row,
            ["현재 모델 (merchant-region adapter)", 0.9123, "강유사 후보", 123.5, "OK"],
        )


if __name__ == "__main__":
    unittest.main()