266 lines
8.2 KiB
Python
266 lines
8.2 KiB
Python
import hashlib
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
from backend.app.services.privatemuse_artifact_installer import (
|
|
ArtifactInstallError,
|
|
install,
|
|
preflight,
|
|
validate_request,
|
|
)
|
|
|
|
|
|
def request(tmp_path: Path, data: bytes = b"model"):
|
|
source = tmp_path / "source.safetensors"
|
|
source.write_bytes(data)
|
|
return {
|
|
"artifact_type": "lora",
|
|
"source_url": "https://models.example.invalid/model.safetensors",
|
|
"sha256": hashlib.sha256(data).hexdigest(),
|
|
"filename": "candidate.safetensors",
|
|
"source_id": "civitai:model/version",
|
|
"license": "upstream terms verified",
|
|
}
|
|
|
|
|
|
def test_validation_is_fail_closed():
|
|
payload = {
|
|
"artifact_type": "lora",
|
|
"source_url": "http://example.test/x",
|
|
"sha256": "0" * 64,
|
|
"filename": "x.safetensors",
|
|
"source_id": "v1",
|
|
"license": "x",
|
|
}
|
|
with pytest.raises(ArtifactInstallError, match="HTTPS"):
|
|
validate_request(payload)
|
|
payload["source_url"] = "https://example.test/x"
|
|
payload["filename"] = "../x.safetensors"
|
|
with pytest.raises(ArtifactInstallError, match="filename"):
|
|
validate_request(payload)
|
|
|
|
|
|
def test_installer_hash_pins_and_writes_receipt(tmp_path, monkeypatch):
|
|
data = b"verified model bytes"
|
|
payload = request(tmp_path, data)
|
|
|
|
class Response:
|
|
headers = {}
|
|
|
|
def __init__(self):
|
|
self.data = data
|
|
self.headers = {"Content-Length": str(len(self.data))}
|
|
|
|
def __enter__(self):
|
|
return self
|
|
|
|
def __exit__(self, *args):
|
|
pass
|
|
|
|
def geturl(self):
|
|
return payload["source_url"]
|
|
|
|
def read(self, n):
|
|
out = self.data[:n]
|
|
self.data = self.data[n:]
|
|
return out
|
|
|
|
monkeypatch.setattr(
|
|
"urllib.request.urlopen",
|
|
lambda *args, **kwargs: Response(),
|
|
)
|
|
result = install(
|
|
payload,
|
|
model_root=tmp_path / "models",
|
|
receipt_root=tmp_path / "receipts",
|
|
)
|
|
assert result["status"] == "installed"
|
|
assert (
|
|
tmp_path / "models" / "loras" / "candidate.safetensors"
|
|
).read_bytes() == data
|
|
assert (tmp_path / "receipts" / f'{payload["sha256"]}.json').exists()
|
|
again = install(payload, model_root=tmp_path / "models", receipt_root=tmp_path / "receipts")
|
|
assert again["status"] == "already_installed"
|
|
assert again["receipt"] == result["receipt"]
|
|
|
|
|
|
def test_installer_rejects_digest_mismatch(tmp_path, monkeypatch):
|
|
payload = request(tmp_path, b"expected")
|
|
|
|
class Response:
|
|
headers = {}
|
|
data = b"different"
|
|
|
|
def __enter__(self):
|
|
return self
|
|
|
|
def __exit__(self, *args):
|
|
pass
|
|
|
|
def geturl(self):
|
|
return payload["source_url"]
|
|
|
|
def read(self, n):
|
|
out = self.data[:n]
|
|
self.data = self.data[n:]
|
|
return out
|
|
|
|
monkeypatch.setattr(
|
|
"urllib.request.urlopen",
|
|
lambda *args, **kwargs: Response(),
|
|
)
|
|
with pytest.raises(ArtifactInstallError, match="SHA-256"):
|
|
install(
|
|
payload,
|
|
model_root=tmp_path / "models",
|
|
receipt_root=tmp_path / "receipts",
|
|
)
|
|
assert not (
|
|
tmp_path / "models" / "loras" / "candidate.safetensors"
|
|
).exists()
|
|
assert not list((tmp_path / "models" / "loras").glob("*.part"))
|
|
|
|
|
|
def test_preflight_requires_persistent_mount_and_detects_changed_destination(tmp_path, monkeypatch):
|
|
payload = request(tmp_path)
|
|
root = tmp_path / "models"
|
|
(root / "loras").mkdir(parents=True)
|
|
with pytest.raises(ArtifactInstallError, match="persistent mount"):
|
|
preflight(payload, model_root=root)
|
|
monkeypatch.setattr("os.path.ismount", lambda path: path == root)
|
|
first = preflight(payload, model_root=root)
|
|
assert first["already_installed"] is False
|
|
(root / "loras" / payload["filename"]).write_bytes(b"model")
|
|
second = preflight(payload, model_root=root)
|
|
assert second["already_installed"] is True
|
|
assert second["fingerprint"] != first["fingerprint"]
|
|
|
|
|
|
def test_preflight_rejects_unwritable_model_directory(tmp_path, monkeypatch):
|
|
payload = request(tmp_path)
|
|
root = tmp_path / "models"
|
|
(root / "loras").mkdir(parents=True)
|
|
monkeypatch.setattr("os.path.ismount", lambda path: path == root)
|
|
monkeypatch.setattr("os.access", lambda path, mode: False)
|
|
with pytest.raises(ArtifactInstallError, match="not writable"):
|
|
preflight(payload, model_root=root)
|
|
|
|
|
|
def test_installer_never_follows_destination_symlink(tmp_path):
|
|
payload = request(tmp_path)
|
|
outside = tmp_path / "outside.safetensors"
|
|
outside.write_bytes(b"model")
|
|
destination_dir = tmp_path / "models" / "loras"
|
|
destination_dir.mkdir(parents=True)
|
|
try:
|
|
(destination_dir / payload["filename"]).symlink_to(outside)
|
|
except OSError:
|
|
pytest.skip("creating symlinks is not permitted on this host")
|
|
with pytest.raises(ArtifactInstallError, match="symlinks"):
|
|
install(payload, model_root=tmp_path / "models", receipt_root=tmp_path / "receipts")
|
|
assert outside.read_bytes() == b"model"
|
|
|
|
|
|
def test_racing_destination_is_never_overwritten(tmp_path, monkeypatch):
|
|
payload = request(tmp_path)
|
|
|
|
class Response:
|
|
headers = {"Content-Length": "5"}
|
|
|
|
def __enter__(self):
|
|
self.data = b"model"
|
|
return self
|
|
|
|
def __exit__(self, *args):
|
|
pass
|
|
|
|
def geturl(self):
|
|
return payload["source_url"]
|
|
|
|
def read(self, size):
|
|
data, self.data = self.data, b""
|
|
return data
|
|
|
|
monkeypatch.setattr("urllib.request.urlopen", lambda *args, **kwargs: Response())
|
|
|
|
def race_link(source, destination):
|
|
Path(destination).write_bytes(b"other")
|
|
raise FileExistsError
|
|
|
|
monkeypatch.setattr("os.link", race_link)
|
|
with pytest.raises(ArtifactInstallError, match="different digest"):
|
|
install(payload, model_root=tmp_path / "models", receipt_root=tmp_path / "receipts")
|
|
assert (tmp_path / "models" / "loras" / payload["filename"]).read_bytes() == b"other"
|
|
assert not list((tmp_path / "models" / "loras").glob("*.part"))
|
|
|
|
|
|
def test_gguf_requires_exact_type_and_extension():
|
|
payload = {
|
|
"artifact_type": "gguf",
|
|
"source_url": "https://models.example.invalid/qwen.gguf",
|
|
"sha256": "a" * 64,
|
|
"filename": "qwen-image-2.1-Q4_K_M.gguf",
|
|
"source_id": "hf:qwen-image-2.1:q4km",
|
|
"license": "qwen-research",
|
|
}
|
|
assert validate_request(payload)["artifact_type"] == "gguf"
|
|
|
|
wrong_type = {**payload, "artifact_type": "lora"}
|
|
with pytest.raises(ArtifactInstallError, match="safetensors"):
|
|
validate_request(wrong_type)
|
|
|
|
wrong_extension = {**payload, "filename": "qwen-image-2.1-Q4_K_M.safetensors"}
|
|
with pytest.raises(ArtifactInstallError, match="gguf"):
|
|
validate_request(wrong_extension)
|
|
|
|
|
|
def test_gguf_installs_only_into_diffusion_models(tmp_path, monkeypatch):
|
|
data = b"gguf model bytes"
|
|
payload = {
|
|
"artifact_type": "gguf",
|
|
"source_url": "https://models.example.invalid/qwen.gguf",
|
|
"sha256": hashlib.sha256(data).hexdigest(),
|
|
"filename": "qwen-image-2.1-Q4_K_M.gguf",
|
|
"source_id": "hf:qwen-image-2.1:q4km",
|
|
"license": "qwen-research",
|
|
}
|
|
|
|
class Response:
|
|
def __init__(self):
|
|
self.data = data
|
|
self.headers = {"Content-Length": str(len(data))}
|
|
|
|
def __enter__(self):
|
|
return self
|
|
|
|
def __exit__(self, *args):
|
|
pass
|
|
|
|
def geturl(self):
|
|
return payload["source_url"]
|
|
|
|
def read(self, n):
|
|
out = self.data[:n]
|
|
self.data = self.data[n:]
|
|
return out
|
|
|
|
monkeypatch.setattr(
|
|
"urllib.request.urlopen",
|
|
lambda *args, **kwargs: Response(),
|
|
)
|
|
result = install(
|
|
payload,
|
|
model_root=tmp_path / "models",
|
|
receipt_root=tmp_path / "receipts",
|
|
)
|
|
expected = (
|
|
tmp_path
|
|
/ "models"
|
|
/ "diffusion_models"
|
|
/ "qwen-image-2.1-Q4_K_M.gguf"
|
|
)
|
|
assert result["status"] == "installed"
|
|
assert expected.read_bytes() == data
|
|
assert not (tmp_path / "models" / "loras" / payload["filename"]).exists()
|