local-translator/tests/test_benchmark.py

import json
from pathlib import Path
import tempfile
import unittest

from local_translator.benchmark import BenchmarkConfig, ModelManager, ModelSpec, aggregate, profile, resolve_metadata, save_report, summarize_run


def item(identity, source, translation, attempts=1, cached=0, status="completed"):
    return {"id": identity, "source": source, "translation": translation, "attempts": attempts, "cached": cached, "status": status}


class MetricsTests(unittest.TestCase):
    def test_count_only_successful_inference_and_keep_missing_usage_unknown(self):
        job = {"id": "job", "status": "completed_with_errors", "items": [
            item("model", "Hello", "Привет", attempts=2), item("local", "Q.LOAD", "Загрузка", attempts=0),
            item("cache", "Yes", "Да", cached=1), item("failed", "Text", None, status="failed"),
        ], "metrics": {"provider_calls": 2, "provider_seconds": 1.5, "completion_tokens": 8, "calls_with_completion_usage": 1}}
        result = summarize_run(job, 2)
        self.assertEqual(result["model_lines"], 1)
        self.assertEqual(result["lines_per_second"], 0.5)
        self.assertEqual(result["chars_per_second"], 2.5)
        self.assertEqual(result["failed"], 1)
        self.assertEqual(result["retry_attempts"], 1)
        self.assertIsNone(result["completion_tokens_per_request_second"])

    def test_reported_token_speed_uses_full_provider_request_time(self):
        job = {"id": "job", "status": "completed", "items": [item("1", "AUTO", "AUTO")],
               "metrics": {"provider_calls": 1, "provider_seconds": 2, "completion_tokens": 40, "calls_with_completion_usage": 1}}
        result = summarize_run(job, 3)
        self.assertEqual(result["completion_tokens_per_request_second"], 20)
        self.assertEqual(result["unchanged_ids"], ["1"])

    def test_failure_is_not_ranked_as_fast_success(self):
        rows = [
            {"model": "A", "batch_size": 10, "status": status, "failed": failed, "seconds": seconds,
             "lines_per_second": rate, "chars_per_second": rate * 5, "retry_attempts": 0, "unchanged_ids": []}
            for status, failed, seconds, rate in [("completed", 0, 10, 1), ("completed", 0, 20, 0.5), ("failed", 10, 0.01, 0)]
        ]
        summary = aggregate(rows)[0]
        self.assertEqual(summary["median_seconds"], 15)
        self.assertEqual(summary["successful_runs"], 2)
        self.assertEqual(summary["runs"], 3)
        self.assertEqual(summary["failed_items"], 10)

    def test_report_escapes_labels_and_preserves_raw_json(self):
        with tempfile.TemporaryDirectory() as directory:
            report = {"dataset": {"source": "<script>alert(1)</script>", "items": 1, "source_chars": 1},
                      "started_at": "test", "runs": [], "errors": ["<b>error</b>"], "restore_errors": []}
            save_report(Path(directory), report)
            document = (Path(directory) / "report.html").read_text(encoding="utf-8")
            self.assertNotIn("<script>", document)
            self.assertIn("&lt;script&gt;", document)
            self.assertEqual(json.loads((Path(directory) / "report.json").read_text(encoding="utf-8"))["errors"], report["errors"])

    def test_native_profiles_and_overrides(self):
        config = BenchmarkConfig()
        hy = profile(ModelSpec(model="tencent/Hy-MT2-7B-GGUF/file.gguf"), config)
        self.assertEqual(hy.prompt_style, "hy_mt")
        self.assertEqual(hy.top_p, 0.6)
        other = profile(ModelSpec(model="other-model", temperature=0.5), config)
        self.assertEqual(other.prompt_style, "generic")
        self.assertEqual(other.temperature, 0.5)
        self.assertIsNone(other.top_k)

    def test_quantizations_require_unambiguous_model_paths(self):
        models = [{"modelKey": "same", "path": "q4.gguf"}, {"modelKey": "same", "path": "q6.gguf"}]
        with self.assertRaises(ValueError):
            resolve_metadata(ModelSpec(model="same"), models)
        self.assertEqual(resolve_metadata(ModelSpec(model="q6.gguf"), models)["path"], "q6.gguf")


class FakeLMS:
    def __init__(self, loaded):
        self.models = [dict(model) for model in loaded]
        self.commands = []

    def json(self, command):
        return [dict(model) for model in self.models]

    def run(self, *args):
        self.commands.append(args)
        if args[0] == "unload":
            self.models = [model for model in self.models if model["identifier"] != args[1]]
        elif args[0] == "load":
            self.models.append({"type": "llm", "modelKey": args[1], "path": args[1], "indexedModelIdentifier": args[1],
                                "identifier": args[args.index("--identifier") + 1],
                                "contextLength": int(args[args.index("--context-length") + 1]), "status": "idle", "queued": 0})
        return ""


class ModelManagementTests(unittest.TestCase):
    def test_switch_and_restore_only_llms_not_embedding(self):
        originals = [
            {"type": "llm", "modelKey": "old", "path": "old.gguf", "indexedModelIdentifier": "old.gguf", "identifier": "original", "contextLength": 4096, "status": "idle", "queued": 0},
            {"type": "embedding", "identifier": "embedding"},
        ]
        lms = FakeLMS(originals)
        manager = ModelManager(lms, BenchmarkConfig())
        loaded, seconds = manager.activate({"modelKey": "new.gguf", "path": "new.gguf", "indexedModelIdentifier": "new.gguf"})
        self.assertTrue(loaded["identifier"].startswith("lt-bench-"))
        self.assertEqual(loaded["contextLength"], 8192)
        manager.restore()
        self.assertEqual(manager.restore_errors, [])
        identifiers = {model["identifier"] for model in lms.models}
        self.assertEqual(identifiers, {"embedding", "original"})
        self.assertFalse(any(command == ("unload", "embedding") for command in lms.commands))

    def test_reuse_loaded_matching_model_without_reloading(self):
        original = {"type": "llm", "modelKey": "old", "path": "old.gguf", "identifier": "original", "contextLength": 8192, "status": "idle", "queued": 0}
        lms = FakeLMS([original])
        manager = ModelManager(lms, BenchmarkConfig(reuse_loaded=True))
        loaded, seconds = manager.activate({"modelKey": "old", "path": "old.gguf"})
        manager.restore()
        self.assertEqual(seconds, 0)
        self.assertEqual(loaded["identifier"], "original")
        self.assertEqual(lms.commands, [])

    def test_busy_models_are_never_unloaded(self):
        lms = FakeLMS([{"type": "llm", "identifier": "busy", "status": "generating", "queued": 0}])
        with self.assertRaises(RuntimeError):
            ModelManager(lms, BenchmarkConfig())
        self.assertEqual(lms.commands, [])


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