tests/test_v110_model_prescan.py

from pathlib import Path

from src.model_scan import analyze_model_metadata, normalize_model_id


def root() -> Path:
    return Path(__file__).resolve().parents[1]


def test_normalize_model_id_accepts_url_and_id():
    assert normalize_model_id("https://huggingface.co/org/model-name?x=1") == "org/model-name"
    assert normalize_model_id("org/model-name") == "org/model-name"


def test_model_prescan_safe_candidate():
    result = analyze_model_metadata(
        model_id="org/safe-model",
        pipeline_tag="text-generation",
        library_name="transformers",
        tags=["safetensors"],
        siblings=[{"rfilename": "config.json"}, {"rfilename": "model.safetensors"}, {"rfilename": "README.md"}],
        readme="This model card has useful documentation. " * 30,
    )
    assert result["verdict"] in {"safe", "caution"}
    assert result["metadata"]["has_safetensors"] is True
    assert any("Safetensors" in item for item in result["good_signals"])


def test_model_prescan_risky_custom_code_and_gated():
    result = analyze_model_metadata(
        model_id="org/risky-model",
        pipeline_tag=None,
        library_name=None,
        gated="manual",
        siblings=[{"rfilename": "model.bin"}, {"rfilename": "modeling_custom.py"}],
        config={"auto_map": {"AutoModel": "modeling_custom.Custom"}},
        readme="tiny",
    )
    assert result["verdict"] in {"risky", "unsupported"}
    assert result["metadata"]["has_custom_code_signal"] is True
    assert any("gated" in item.lower() for item in result["risk_signals"])


def test_prescan_route_and_ui_are_present():
    app = (root() / "app.py").read_text(encoding="utf-8")
    html = (root() / "web" / "index.html").read_text(encoding="utf-8")
    js = (root() / "web" / "static" / "app.js").read_text(encoding="utf-8")
    css = (root() / "web" / "static" / "app.css").read_text(encoding="utf-8")
    assert '"/api/models/pre-scan"' in app
    assert "modelPreScanCard" in html
    assert "scanModel" in js
    assert "modelScanMatchesCurrent" in js
    assert "Run the model pre-scan before launching" in js
    assert ".model-prescan-card" in css


def test_model_prescan_known_good_z_image_turbo_is_safe():
    readme = """
    # Z-Image-Turbo

    This is a text-to-image Diffusers model.

    ```python
    import torch
    from diffusers import DiffusionPipeline

    pipe = DiffusionPipeline.from_pretrained(
        "Tongyi-MAI/Z-Image-Turbo",
        torch_dtype=torch.bfloat16,
        device_map="cuda",
    )
    image = pipe(prompt="A small robot building a Hugging Face Space", num_inference_steps=8).images[0]
    ```
    """ * 5
    result = analyze_model_metadata(
        model_id="Tongyi-MAI/Z-Image-Turbo",
        pipeline_tag="text-to-image",
        library_name="diffusers",
        tags=["diffusers", "safetensors", "text-to-image"],
        siblings=[
            {"rfilename": "README.md"},
            {"rfilename": "model_index.json"},
            {"rfilename": "transformer/diffusion_pytorch_model.safetensors"},
            {"rfilename": "vae/diffusion_pytorch_model.safetensors"},
            {"rfilename": "scheduler/scheduler_config.json"},
        ],
        model_index={"_class_name": "ZImagePipeline"},
        readme=readme,
    )
    assert result["verdict"] == "safe"
    assert result["score"] >= 82
    assert result["metadata"]["has_diffusers_example"] is True
    assert result["metadata"]["diffusers_standard"] is True
    assert result["metadata"]["pipeline_class"] == "ZImagePipeline"
    assert any("Diffusers example" in item for item in result["good_signals"])


def test_default_model_is_known_good_z_image_turbo():
    html = (root() / "web" / "index.html").read_text(encoding="utf-8")
    assert 'value="Tongyi-MAI/Z-Image-Turbo"' in html
    assert 'placeholder="z-image-turbo-demo"' in html