mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-03-22 05:32:12 -03:00
46 lines
1.1 KiB
Python
46 lines
1.1 KiB
Python
import pytest
|
|
|
|
from py.utils import lora_metadata
|
|
|
|
|
|
class DummySafeOpen:
|
|
def __init__(self, metadata):
|
|
self._metadata = metadata
|
|
|
|
def __enter__(self):
|
|
return self
|
|
|
|
def __exit__(self, *args):
|
|
return False
|
|
|
|
def metadata(self):
|
|
return self._metadata
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_extract_lora_metadata_returns_base_model(monkeypatch, tmp_path):
|
|
file_path = tmp_path / "model.safetensors"
|
|
file_path.write_bytes(b"")
|
|
|
|
monkeypatch.setattr(
|
|
lora_metadata,
|
|
"safe_open",
|
|
lambda *args, **kwargs: DummySafeOpen({"ss_base_model_version": "sdxl"}),
|
|
)
|
|
|
|
metadata = await lora_metadata.extract_lora_metadata(str(file_path))
|
|
|
|
assert metadata == {"base_model": "SDXL 1.0"}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_extract_lora_metadata_handles_errors(monkeypatch):
|
|
def raising_safe_open(*_, **__):
|
|
raise RuntimeError("boom")
|
|
|
|
monkeypatch.setattr(lora_metadata, "safe_open", raising_safe_open)
|
|
|
|
metadata = await lora_metadata.extract_lora_metadata("missing.safetensors")
|
|
|
|
assert metadata == {"base_model": "Unknown"}
|