benchmarks/longmemeval/tests/test_safe_model_loading.py

from __future__ import annotations

import ast
import importlib.util
import json
from pathlib import Path

import pytest


ROOT = Path(__file__).resolve().parents[1]
SOURCE_ROOT = ROOT / "src" / "tmcra_benchmark"
REPOSITORY_ROOT = ROOT.parents[1]
SECURITY_SOURCE_ROOTS = (SOURCE_ROOT, REPOSITORY_ROOT / "code")


def _dotted_name(node: ast.AST) -> str:
    if isinstance(node, ast.Name):
        return node.id
    if isinstance(node, ast.Attribute):
        prefix = _dotted_name(node.value)
        return f"{prefix}.{node.attr}" if prefix else node.attr
    return ""


def _keyword(call: ast.Call, name: str) -> ast.AST | None:
    return next((item.value for item in call.keywords if item.arg == name), None)


def test_every_torch_load_is_weights_only() -> None:
    violations: list[str] = []
    for source_root in SECURITY_SOURCE_ROOTS:
        for path in source_root.rglob("*.py"):
            tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path))
            torch_aliases = {
                alias.asname or alias.name
                for node in ast.walk(tree)
                if isinstance(node, ast.Import)
                for alias in node.names
                if alias.name == "torch"
            }
            torch_aliases.update({"torch", "torch_module"})
            for node in ast.walk(tree):
                if not isinstance(node, ast.Call) or not isinstance(node.func, ast.Attribute):
                    continue
                if node.func.attr != "load":
                    continue
                receiver = _dotted_name(node.func.value)
                if receiver not in torch_aliases and not receiver.endswith(".torch"):
                    continue
                value = _keyword(node, "weights_only")
                if not isinstance(value, ast.Constant) or value.value is not True:
                    violations.append(f"{path.relative_to(REPOSITORY_ROOT)}:{node.lineno}")
    assert not violations, "torch.load must use weights_only=True: " + ", ".join(violations)


def test_local_transformer_models_disable_remote_code() -> None:
    violations: list[str] = []
    for path in SOURCE_ROOT.rglob("*.py"):
        tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path))
        for node in ast.walk(tree):
            if not isinstance(node, ast.Call) or not isinstance(node.func, ast.Attribute):
                continue
            if node.func.attr != "from_pretrained":
                continue
            owner = _dotted_name(node.func.value)
            if not (owner.startswith("AutoTokenizer") or owner.startswith("AutoModel")):
                continue
            local_only = _keyword(node, "local_files_only")
            remote_code = _keyword(node, "trust_remote_code")
            if not isinstance(local_only, ast.Constant) or local_only.value is not True:
                violations.append(f"{path.relative_to(ROOT)}:{node.lineno}:local_files_only")
            if not isinstance(remote_code, ast.Constant) or remote_code.value is not False:
                violations.append(f"{path.relative_to(ROOT)}:{node.lineno}:trust_remote_code")
    assert not violations, "local model loading is not hardened: " + ", ".join(violations)


def test_public_code_never_enables_remote_code() -> None:
    violations: list[str] = []
    for source_root in SECURITY_SOURCE_ROOTS:
        for path in source_root.rglob("*.py"):
            tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path))
            for node in ast.walk(tree):
                if not isinstance(node, ast.Call) or node.func is None:
                    continue
                remote_code = _keyword(node, "trust_remote_code")
                if isinstance(remote_code, ast.Constant) and remote_code.value is True:
                    violations.append(f"{path.relative_to(REPOSITORY_ROOT)}:{node.lineno}")
    assert not violations, "public code enables trust_remote_code: " + ", ".join(violations)


def _load_model_assets_module():
    path = SOURCE_ROOT / "legacy" / "tmcra_model_assets.py"
    spec = importlib.util.spec_from_file_location(path.stem, path)
    assert spec and spec.loader
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module


def test_official_hf_model_does_not_require_tmcra_manifest(tmp_path: Path) -> None:
    module = _load_model_assets_module()
    (tmp_path / "config.json").write_text("{}\n", encoding="utf-8")

    assert module.load_optional_hf_model_manifest(tmp_path) == {}

    for relative in ("tmcra_v3_reranker.py", "tmcra_v3_online_runtime.py"):
        source = (SOURCE_ROOT / "legacy" / relative).read_text(encoding="utf-8")
        assert "load_optional_hf_model_manifest" in source
        assert "pinned cross model manifest is required" not in source
        assert "pinned model manifest is required" not in source


def test_optional_tmcra_manifest_is_validated(tmp_path: Path) -> None:
    module = _load_model_assets_module()
    manifest = tmp_path / "TMCRA_MODEL_MANIFEST.json"
    manifest.write_text(json.dumps({"repo_id": "BAAI/bge-reranker-v2-m3", "revision": "abc"}), encoding="utf-8")
    assert module.load_optional_hf_model_manifest(tmp_path)["revision"] == "abc"

    manifest.write_text("[]", encoding="utf-8")
    with pytest.raises(ValueError, match="JSON object"):
        module.load_optional_hf_model_manifest(tmp_path)