harness_zbt/backend/tests/test_bge_model_loader.py

from pathlib import Path
from types import SimpleNamespace

from zbt.infrastructure.embedding.bge_model import (
    clear_bge_m3_model_cache,
    load_bge_m3_model,
    resolve_local_huggingface_snapshot,
)


def _create_complete_snapshot(cache_root: Path, revision: str = "revision-01") -> Path:
    repository = cache_root / "models--BAAI--bge-m3"
    snapshot = repository / "snapshots" / revision
    snapshot.mkdir(parents=True)
    (repository / "refs").mkdir()
    (repository / "refs" / "main").write_text(revision, encoding="utf-8")
    for filename in (
        "config.json",
        "tokenizer_config.json",
        "tokenizer.json",
        "model.safetensors",
    ):
        (snapshot / filename).write_text("fixture", encoding="utf-8")
    return snapshot


def test_resolve_local_huggingface_snapshot_uses_complete_cached_model(
    tmp_path: Path,
) -> None:
    snapshot = _create_complete_snapshot(tmp_path)

    result = resolve_local_huggingface_snapshot("BAAI/bge-m3", cache_dir=tmp_path)

    assert result == str(snapshot.resolve())


def test_resolve_local_huggingface_snapshot_keeps_remote_id_when_cache_incomplete(
    tmp_path: Path,
) -> None:
    repository = tmp_path / "models--BAAI--bge-m3" / "snapshots" / "incomplete"
    repository.mkdir(parents=True)
    (repository / "config.json").write_text("fixture", encoding="utf-8")

    result = resolve_local_huggingface_snapshot("BAAI/bge-m3", cache_dir=tmp_path)

    assert result == "BAAI/bge-m3"


def test_load_bge_m3_model_reuses_one_process_instance(
    monkeypatch,
    tmp_path: Path,
) -> None:
    snapshot = _create_complete_snapshot(tmp_path)
    created_with: list[tuple[str, str]] = []

    class FakeBgeM3FlagModel:
        def __init__(self, model_path: str, **options: object) -> None:
            created_with.append((model_path, str(options["devices"])))

    monkeypatch.setitem(
        __import__("sys").modules,
        "FlagEmbedding",
        SimpleNamespace(BGEM3FlagModel=FakeBgeM3FlagModel),
    )
    clear_bge_m3_model_cache()

    first = load_bge_m3_model(
        "BAAI/bge-m3",
        "cpu",
        cache_dir=tmp_path,
    )
    second = load_bge_m3_model(
        "BAAI/bge-m3",
        "cpu",
        cache_dir=tmp_path,
    )

    assert first is second
    assert created_with == [(str(snapshot.resolve()), "cpu")]
    clear_bge_m3_model_cache()