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