415 lines
13 KiB
Python
415 lines
13 KiB
Python
"""Tests for zero-touch Hugging Face model provenance discovery."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import struct
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
|
|
class FakeHubApi:
|
|
"""Return deterministic file metadata for candidate verification tests."""
|
|
|
|
def __init__(self, responses: dict[tuple[str, str], object]) -> None:
|
|
"""Store responses indexed by repository and requested revision."""
|
|
self.responses = responses
|
|
self.calls: list[tuple[str, str]] = []
|
|
|
|
def model_info(self, repo_id: str, **kwargs: object) -> object:
|
|
"""Return one configured repository response and record the request."""
|
|
revision = str(kwargs["revision"])
|
|
assert kwargs["files_metadata"] is True
|
|
self.calls.append((repo_id, revision))
|
|
return self.responses[(repo_id, revision)]
|
|
|
|
|
|
def _hub_response(
|
|
*,
|
|
filename: str,
|
|
sha256: str,
|
|
size_bytes: int,
|
|
commit: str,
|
|
) -> dict[str, object]:
|
|
"""Build one Hub model-info payload containing exact file metadata."""
|
|
return {
|
|
"sha": commit,
|
|
"siblings": [
|
|
{
|
|
"rfilename": filename,
|
|
"size": size_bytes,
|
|
"lfs": {"sha256": sha256},
|
|
}
|
|
],
|
|
}
|
|
|
|
|
|
def _discovery(
|
|
module: Any,
|
|
assets_module: Any,
|
|
*,
|
|
tmp_path: Path,
|
|
user_directory: Path,
|
|
api: FakeHubApi,
|
|
where_from_urls: tuple[str, ...] = (),
|
|
) -> Any:
|
|
"""Create an isolated discovery service with deterministic local evidence."""
|
|
registry = assets_module.HuggingFaceAssetRegistry(
|
|
tmp_path / "registry" / "huggingface-assets.json"
|
|
)
|
|
return module.HuggingFaceAssetDiscovery(
|
|
registry=registry,
|
|
user_directory=user_directory,
|
|
comfyui_root=tmp_path / "comfyui",
|
|
api=api,
|
|
where_from_reader=lambda path: where_from_urls,
|
|
)
|
|
|
|
|
|
def test_parses_official_resolve_and_cache_urls(
|
|
huggingface_discovery_module: Any,
|
|
) -> None:
|
|
"""Normal and CDN-backed Hub URLs should yield the same immutable hint shape."""
|
|
resolve_hint = huggingface_discovery_module.huggingface_hint_from_url(
|
|
"https://huggingface.co/owner/model/resolve/main/weights/model.safetensors?download=true",
|
|
evidence="browser",
|
|
)
|
|
cache_hint = huggingface_discovery_module.huggingface_hint_from_url(
|
|
"https://huggingface.co/api/resolve-cache/models/owner/model/"
|
|
f"{'a' * 40}/weights/model.safetensors",
|
|
evidence="browser",
|
|
)
|
|
|
|
assert resolve_hint.repo_id == "owner/model"
|
|
assert resolve_hint.revision == "main"
|
|
assert resolve_hint.filename == "weights/model.safetensors"
|
|
assert cache_hint.repo_id == "owner/model"
|
|
assert cache_hint.revision == "a" * 40
|
|
assert cache_hint.filename == "weights/model.safetensors"
|
|
|
|
|
|
def test_discovers_manager_installed_model_and_persists_verified_mapping(
|
|
huggingface_assets_module: Any,
|
|
huggingface_discovery_module: Any,
|
|
tmp_path: Path,
|
|
) -> None:
|
|
"""Manager's catalog should make ordinary installs zero-touch after hash validation."""
|
|
user_directory = tmp_path / "user"
|
|
catalog = user_directory / "__manager" / "cache" / "123_model-list.json"
|
|
catalog.parent.mkdir(parents=True)
|
|
catalog.write_text(
|
|
json.dumps(
|
|
{
|
|
"models": [
|
|
{
|
|
"filename": "manager-model.safetensors",
|
|
"save_path": "vae",
|
|
"url": (
|
|
"https://huggingface.co/wrong/model/resolve/main/"
|
|
"manager-model.safetensors"
|
|
),
|
|
},
|
|
{
|
|
"filename": "manager-model.safetensors",
|
|
"save_path": "checkpoints",
|
|
"url": (
|
|
"https://huggingface.co/owner/model/resolve/main/"
|
|
"weights/manager-model.safetensors"
|
|
),
|
|
}
|
|
]
|
|
}
|
|
),
|
|
encoding="utf-8",
|
|
)
|
|
asset_path = tmp_path / "models" / "checkpoints" / "manager-model.safetensors"
|
|
asset_path.parent.mkdir(parents=True)
|
|
asset_path.write_bytes(b"manager-model")
|
|
sha256 = huggingface_assets_module.sha256_file(asset_path)
|
|
commit = "b" * 40
|
|
api = FakeHubApi(
|
|
{
|
|
("owner/model", "main"): _hub_response(
|
|
filename="weights/manager-model.safetensors",
|
|
sha256=sha256,
|
|
size_bytes=asset_path.stat().st_size,
|
|
commit=commit,
|
|
)
|
|
}
|
|
)
|
|
discovery = _discovery(
|
|
huggingface_discovery_module,
|
|
huggingface_assets_module,
|
|
tmp_path=tmp_path,
|
|
user_directory=user_directory,
|
|
api=api,
|
|
)
|
|
|
|
source = discovery.discover(asset_path, sha256=sha256)
|
|
|
|
assert source.repo_id == "owner/model"
|
|
assert source.revision == commit
|
|
assert source.filename == "weights/manager-model.safetensors"
|
|
assert discovery.registry.get(sha256) == source
|
|
assert discovery.discover(asset_path, sha256=sha256) == source
|
|
assert api.calls == [("owner/model", "main")]
|
|
|
|
|
|
def test_discovers_model_inside_huggingface_git_checkout(
|
|
huggingface_assets_module: Any,
|
|
huggingface_discovery_module: Any,
|
|
tmp_path: Path,
|
|
) -> None:
|
|
"""An ordinary Hub Git clone should supply its repository, commit, and file path."""
|
|
commit = "8" * 40
|
|
checkout = tmp_path / "checkout"
|
|
git_directory = checkout / ".git"
|
|
git_directory.mkdir(parents=True)
|
|
(git_directory / "config").write_text(
|
|
'[remote "origin"]\n\turl = https://huggingface.co/git/model.git\n',
|
|
encoding="utf-8",
|
|
)
|
|
(git_directory / "HEAD").write_text(f"{commit}\n", encoding="utf-8")
|
|
asset_path = checkout / "weights" / "git.safetensors"
|
|
asset_path.parent.mkdir()
|
|
asset_path.write_bytes(b"git-model")
|
|
sha256 = huggingface_assets_module.sha256_file(asset_path)
|
|
api = FakeHubApi(
|
|
{
|
|
("git/model", commit): _hub_response(
|
|
filename="weights/git.safetensors",
|
|
sha256=sha256,
|
|
size_bytes=asset_path.stat().st_size,
|
|
commit=commit,
|
|
)
|
|
}
|
|
)
|
|
discovery = _discovery(
|
|
huggingface_discovery_module,
|
|
huggingface_assets_module,
|
|
tmp_path=tmp_path,
|
|
user_directory=tmp_path / "user",
|
|
api=api,
|
|
)
|
|
|
|
source = discovery.discover(asset_path, sha256=sha256)
|
|
|
|
assert source.repo_id == "git/model"
|
|
assert source.revision == commit
|
|
assert source.filename == "weights/git.safetensors"
|
|
|
|
|
|
def test_discovers_browser_download_source_metadata(
|
|
huggingface_assets_module: Any,
|
|
huggingface_discovery_module: Any,
|
|
tmp_path: Path,
|
|
) -> None:
|
|
"""Browser WhereFrom metadata should remove the need for manual registration."""
|
|
asset_path = tmp_path / "browser-model.safetensors"
|
|
asset_path.write_bytes(b"browser-model")
|
|
sha256 = huggingface_assets_module.sha256_file(asset_path)
|
|
api = FakeHubApi(
|
|
{
|
|
("browser/source", "main"): _hub_response(
|
|
filename="browser-model.safetensors",
|
|
sha256=sha256,
|
|
size_bytes=asset_path.stat().st_size,
|
|
commit="c" * 40,
|
|
)
|
|
}
|
|
)
|
|
discovery = _discovery(
|
|
huggingface_discovery_module,
|
|
huggingface_assets_module,
|
|
tmp_path=tmp_path,
|
|
user_directory=tmp_path / "user",
|
|
api=api,
|
|
where_from_urls=(
|
|
"https://huggingface.co/browser/source/resolve/main/browser-model.safetensors",
|
|
),
|
|
)
|
|
|
|
source = discovery.discover(asset_path, sha256=sha256)
|
|
|
|
assert source.repo_id == "browser/source"
|
|
assert source.revision == "c" * 40
|
|
|
|
|
|
def test_discovers_standard_huggingface_cache_path(
|
|
huggingface_assets_module: Any,
|
|
huggingface_discovery_module: Any,
|
|
tmp_path: Path,
|
|
) -> None:
|
|
"""A file or symlink into the standard Hub cache should carry its own source."""
|
|
commit = "d" * 40
|
|
asset_path = (
|
|
tmp_path
|
|
/ "hub"
|
|
/ "models--cache--model"
|
|
/ "snapshots"
|
|
/ commit
|
|
/ "weights"
|
|
/ "cached.safetensors"
|
|
)
|
|
asset_path.parent.mkdir(parents=True)
|
|
asset_path.write_bytes(b"cached-model")
|
|
sha256 = huggingface_assets_module.sha256_file(asset_path)
|
|
api = FakeHubApi(
|
|
{
|
|
("cache/model", commit): _hub_response(
|
|
filename="weights/cached.safetensors",
|
|
sha256=sha256,
|
|
size_bytes=asset_path.stat().st_size,
|
|
commit=commit,
|
|
)
|
|
}
|
|
)
|
|
discovery = _discovery(
|
|
huggingface_discovery_module,
|
|
huggingface_assets_module,
|
|
tmp_path=tmp_path,
|
|
user_directory=tmp_path / "user",
|
|
api=api,
|
|
)
|
|
|
|
source = discovery.discover(asset_path, sha256=sha256)
|
|
|
|
assert source.repo_id == "cache/model"
|
|
assert source.filename == "weights/cached.safetensors"
|
|
|
|
|
|
def test_discovers_hub_url_embedded_in_safetensors_metadata(
|
|
huggingface_assets_module: Any,
|
|
huggingface_discovery_module: Any,
|
|
tmp_path: Path,
|
|
) -> None:
|
|
"""Model metadata containing a Hub file URL should be useful but still verified."""
|
|
header = json.dumps(
|
|
{
|
|
"__metadata__": {
|
|
"modelspec.source": (
|
|
"https://huggingface.co/embedded/model/resolve/main/"
|
|
"weights/embedded.safetensors"
|
|
)
|
|
}
|
|
}
|
|
).encode("utf-8")
|
|
asset_path = tmp_path / "embedded.safetensors"
|
|
asset_path.write_bytes(struct.pack("<Q", len(header)) + header + b"model-bytes")
|
|
sha256 = huggingface_assets_module.sha256_file(asset_path)
|
|
api = FakeHubApi(
|
|
{
|
|
("embedded/model", "main"): _hub_response(
|
|
filename="weights/embedded.safetensors",
|
|
sha256=sha256,
|
|
size_bytes=asset_path.stat().st_size,
|
|
commit="e" * 40,
|
|
)
|
|
}
|
|
)
|
|
discovery = _discovery(
|
|
huggingface_discovery_module,
|
|
huggingface_assets_module,
|
|
tmp_path=tmp_path,
|
|
user_directory=tmp_path / "user",
|
|
api=api,
|
|
)
|
|
|
|
source = discovery.discover(asset_path, sha256=sha256)
|
|
|
|
assert source.repo_id == "embedded/model"
|
|
assert source.filename == "weights/embedded.safetensors"
|
|
|
|
|
|
def test_rejects_unverified_hint_and_negative_caches_result(
|
|
huggingface_assets_module: Any,
|
|
huggingface_discovery_module: Any,
|
|
tmp_path: Path,
|
|
) -> None:
|
|
"""Install hints must never bypass exact content verification or repeat per prompt."""
|
|
asset_path = tmp_path / "mismatch.safetensors"
|
|
asset_path.write_bytes(b"local-content")
|
|
sha256 = huggingface_assets_module.sha256_file(asset_path)
|
|
api = FakeHubApi(
|
|
{
|
|
("wrong/model", "main"): _hub_response(
|
|
filename="mismatch.safetensors",
|
|
sha256="f" * 64,
|
|
size_bytes=asset_path.stat().st_size,
|
|
commit="f" * 40,
|
|
)
|
|
}
|
|
)
|
|
discovery = _discovery(
|
|
huggingface_discovery_module,
|
|
huggingface_assets_module,
|
|
tmp_path=tmp_path,
|
|
user_directory=tmp_path / "user",
|
|
api=api,
|
|
where_from_urls=(
|
|
"https://huggingface.co/wrong/model/resolve/main/mismatch.safetensors",
|
|
),
|
|
)
|
|
|
|
assert discovery.discover(asset_path, sha256=sha256) is None
|
|
assert discovery.discover(asset_path, sha256=sha256) is None
|
|
assert discovery.registry.get(sha256) is None
|
|
assert api.calls == [("wrong/model", "main")]
|
|
|
|
|
|
def test_retries_unresolved_asset_after_manager_catalog_changes(
|
|
huggingface_assets_module: Any,
|
|
huggingface_discovery_module: Any,
|
|
tmp_path: Path,
|
|
) -> None:
|
|
"""Installing a model while ComfyUI is open should invalidate negative evidence."""
|
|
user_directory = tmp_path / "user"
|
|
asset_path = tmp_path / "models" / "late.safetensors"
|
|
asset_path.parent.mkdir()
|
|
asset_path.write_bytes(b"late-manager-install")
|
|
sha256 = huggingface_assets_module.sha256_file(asset_path)
|
|
api = FakeHubApi(
|
|
{
|
|
("late/model", "main"): _hub_response(
|
|
filename="late.safetensors",
|
|
sha256=sha256,
|
|
size_bytes=asset_path.stat().st_size,
|
|
commit="7" * 40,
|
|
)
|
|
}
|
|
)
|
|
discovery = _discovery(
|
|
huggingface_discovery_module,
|
|
huggingface_assets_module,
|
|
tmp_path=tmp_path,
|
|
user_directory=user_directory,
|
|
api=api,
|
|
)
|
|
|
|
assert discovery.discover(asset_path, sha256=sha256) is None
|
|
|
|
catalog = user_directory / "__manager" / "cache" / "new_model-list.json"
|
|
catalog.parent.mkdir(parents=True)
|
|
catalog.write_text(
|
|
json.dumps(
|
|
{
|
|
"models": [
|
|
{
|
|
"filename": "late.safetensors",
|
|
"url": (
|
|
"https://huggingface.co/late/model/resolve/main/"
|
|
"late.safetensors"
|
|
),
|
|
}
|
|
]
|
|
}
|
|
),
|
|
encoding="utf-8",
|
|
)
|
|
|
|
source = discovery.discover(asset_path, sha256=sha256)
|
|
|
|
assert source.repo_id == "late/model"
|
|
assert source.revision == "7" * 40
|
|
assert api.calls == [("late/model", "main")]
|