436 lines
14 KiB
Python
436 lines
14 KiB
Python
import base64
|
|
import json
|
|
import os
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import pytest
|
|
from aiohttp import web
|
|
from src.nodes.style_loader import STYLE_TEMPLATE, _apply_loras
|
|
from src.utils_loader import safe_relative_path
|
|
|
|
# --- _apply_loras ---
|
|
|
|
|
|
def make_model_clip():
|
|
return MagicMock(), MagicMock()
|
|
|
|
|
|
def test_apply_loras_empty():
|
|
model, clip = make_model_clip()
|
|
result_model, result_clip = _apply_loras(model, clip, [])
|
|
assert result_model is model
|
|
assert result_clip is clip
|
|
|
|
|
|
def test_apply_loras_skips_missing(capsys):
|
|
model, clip = make_model_clip()
|
|
with patch("folder_paths.get_full_path", return_value=None):
|
|
result_model, result_clip = _apply_loras(
|
|
model, clip, [{"name": "missing.safetensors"}]
|
|
)
|
|
captured = capsys.readouterr()
|
|
assert "not found" in captured.out
|
|
assert result_model is model
|
|
assert result_clip is clip
|
|
|
|
|
|
def test_apply_loras_skips_empty_name():
|
|
model, clip = make_model_clip()
|
|
result_model, result_clip = _apply_loras(model, clip, [{"name": ""}])
|
|
assert result_model is model
|
|
assert result_clip is clip
|
|
|
|
|
|
def test_apply_loras_weight_fallback():
|
|
model, clip = make_model_clip()
|
|
fake_lora_sd = {}
|
|
with (
|
|
patch("folder_paths.get_full_path", return_value="/fake/lora.safetensors"),
|
|
patch("comfy.utils.load_torch_file", return_value=fake_lora_sd),
|
|
patch("comfy.sd.load_lora_for_models", return_value=(model, clip)) as mock_load,
|
|
):
|
|
_apply_loras(model, clip, [{"name": "lora.safetensors", "weight": 0.8}])
|
|
mock_load.assert_called_once_with(model, clip, fake_lora_sd, 0.8, 0.8)
|
|
|
|
|
|
def test_apply_loras_explicit_weights():
|
|
model, clip = make_model_clip()
|
|
fake_lora_sd = {}
|
|
with (
|
|
patch("folder_paths.get_full_path", return_value="/fake/lora.safetensors"),
|
|
patch("comfy.utils.load_torch_file", return_value=fake_lora_sd),
|
|
patch("comfy.sd.load_lora_for_models", return_value=(model, clip)) as mock_load,
|
|
):
|
|
_apply_loras(
|
|
model,
|
|
clip,
|
|
[{"name": "lora.safetensors", "model_weight": 0.7, "clip_weight": 0.4}],
|
|
)
|
|
mock_load.assert_called_once_with(model, clip, fake_lora_sd, 0.7, 0.4)
|
|
|
|
|
|
def test_apply_loras_chained():
|
|
model1, clip1 = make_model_clip()
|
|
model2, clip2 = make_model_clip()
|
|
model3, clip3 = make_model_clip()
|
|
fake_lora_sd = {}
|
|
call_results = [(model2, clip2), (model3, clip3)]
|
|
with (
|
|
patch("folder_paths.get_full_path", return_value="/fake/lora.safetensors"),
|
|
patch("comfy.utils.load_torch_file", return_value=fake_lora_sd),
|
|
patch("comfy.sd.load_lora_for_models", side_effect=call_results),
|
|
):
|
|
result_model, result_clip = _apply_loras(
|
|
model1,
|
|
clip1,
|
|
[
|
|
{"name": "lora1.safetensors", "weight": 0.8},
|
|
{"name": "lora2.safetensors", "weight": 0.6},
|
|
],
|
|
)
|
|
assert result_model is model3
|
|
assert result_clip is clip3
|
|
|
|
|
|
# --- STYLE_TEMPLATE ---
|
|
|
|
|
|
def test_style_template_has_required_keys():
|
|
for key in [
|
|
"checkpoint",
|
|
"vae",
|
|
"clip_skip",
|
|
"loras",
|
|
"quality_tags",
|
|
"extra_quality_tags",
|
|
"negative_tags",
|
|
"steps",
|
|
"refiner_step",
|
|
"cfg",
|
|
]:
|
|
assert key in STYLE_TEMPLATE
|
|
|
|
|
|
def test_style_template_loras_empty_list():
|
|
assert STYLE_TEMPLATE["loras"] == []
|
|
|
|
|
|
def test_style_template_defaults():
|
|
assert STYLE_TEMPLATE["clip_skip"] == -2
|
|
assert STYLE_TEMPLATE["cfg"] == 4.0
|
|
assert STYLE_TEMPLATE["steps"] == 15
|
|
|
|
|
|
# --- Routes ---
|
|
|
|
|
|
@pytest.fixture
|
|
def app(tmp_path):
|
|
async def save_handler(request):
|
|
body = await request.json()
|
|
file = body.get("file", "")
|
|
content = body.get("content", "")
|
|
rel = safe_relative_path(file)
|
|
if not rel:
|
|
return web.json_response({"error": "Invalid file path."}, status=400)
|
|
try:
|
|
json.loads(content)
|
|
except json.JSONDecodeError as e:
|
|
return web.json_response({"error": f"Invalid JSON: {e}"}, status=400)
|
|
path = os.path.join(str(tmp_path), rel)
|
|
os.makedirs(os.path.dirname(path), exist_ok=True)
|
|
with open(path, "w", encoding="utf-8") as f:
|
|
f.write(content)
|
|
return web.json_response({"ok": True})
|
|
|
|
async def new_handler(request):
|
|
body = await request.json()
|
|
file = body.get("file", "").strip()
|
|
if not file:
|
|
return web.json_response({"error": "File name is required."}, status=400)
|
|
if not file.endswith(".json"):
|
|
file += ".json"
|
|
rel = safe_relative_path(file)
|
|
if not rel:
|
|
return web.json_response({"error": "Invalid file path."}, status=400)
|
|
path = os.path.join(str(tmp_path), rel)
|
|
if os.path.exists(path):
|
|
return web.json_response({"error": "File already exists."}, status=409)
|
|
os.makedirs(os.path.dirname(path), exist_ok=True)
|
|
content = json.dumps(STYLE_TEMPLATE, indent=2, ensure_ascii=False)
|
|
with open(path, "w", encoding="utf-8") as f:
|
|
f.write(content)
|
|
return web.json_response(
|
|
{"ok": True, "file": rel.replace("\\", "/"), "content": content}
|
|
)
|
|
|
|
async def load_handler(request):
|
|
file = request.rel_url.query.get("file", "")
|
|
rel = safe_relative_path(file)
|
|
if not rel:
|
|
return web.json_response({"error": "Invalid file path."}, status=400)
|
|
path = os.path.join(str(tmp_path), rel)
|
|
try:
|
|
with open(path, "r", encoding="utf-8") as f:
|
|
content = f.read()
|
|
return web.json_response({"content": content})
|
|
except FileNotFoundError:
|
|
return web.json_response({"error": "File not found."}, status=404)
|
|
|
|
async def delete_handler(request):
|
|
body = await request.json()
|
|
file = body.get("file", "")
|
|
rel = safe_relative_path(file)
|
|
if not rel:
|
|
return web.json_response({"error": "Invalid file path."}, status=400)
|
|
path = os.path.join(str(tmp_path), rel)
|
|
if not os.path.isfile(path):
|
|
return web.json_response({"error": "File not found."}, status=404)
|
|
os.remove(path)
|
|
return web.json_response({"ok": True})
|
|
|
|
async def preview_handler(request):
|
|
file = request.rel_url.query.get("file", "")
|
|
rel = safe_relative_path(file)
|
|
if not rel:
|
|
return web.json_response({"error": "Invalid path."}, status=400)
|
|
|
|
base = os.path.splitext(os.path.join(str(tmp_path), rel))[0]
|
|
for ext in [".png", ".jpg", ".webp", ".jpeg"]:
|
|
path = base + ext
|
|
if os.path.isfile(path):
|
|
with open(path, "rb") as f:
|
|
data = f.read()
|
|
b64 = base64.b64encode(data).decode()
|
|
mime = (
|
|
"image/png"
|
|
if ext == ".png"
|
|
else "image/jpeg"
|
|
if ext in [".jpg", ".jpeg"]
|
|
else "image/webp"
|
|
)
|
|
return web.json_response({"image": f"data:{mime};base64,{b64}"})
|
|
|
|
return web.json_response({"image": None})
|
|
|
|
application = web.Application()
|
|
application.router.add_post("/styles/save", save_handler)
|
|
application.router.add_post("/styles/new", new_handler)
|
|
application.router.add_get("/styles/load", load_handler)
|
|
application.router.add_post("/styles/delete", delete_handler)
|
|
application.router.add_get("/styles/preview", preview_handler)
|
|
return application
|
|
|
|
|
|
@pytest.fixture
|
|
async def client(app, aiohttp_client):
|
|
return await aiohttp_client(app)
|
|
|
|
|
|
# --- /styles/save ---
|
|
|
|
|
|
async def test_save_valid(client, tmp_path):
|
|
content = json.dumps({"checkpoint": "model.safetensors"})
|
|
resp = await client.post(
|
|
"/styles/save", json={"file": "test.json", "content": content}
|
|
)
|
|
assert resp.status == 200
|
|
data = await resp.json()
|
|
assert data["ok"] is True
|
|
assert (tmp_path / "test.json").exists()
|
|
|
|
|
|
async def test_save_invalid_json(client):
|
|
resp = await client.post(
|
|
"/styles/save", json={"file": "test.json", "content": "not json"}
|
|
)
|
|
assert resp.status == 400
|
|
data = await resp.json()
|
|
assert "Invalid JSON" in data["error"]
|
|
|
|
|
|
async def test_save_invalid_path(client):
|
|
resp = await client.post(
|
|
"/styles/save", json={"file": "../../evil.json", "content": "{}"}
|
|
)
|
|
assert resp.status == 400
|
|
|
|
|
|
async def test_save_creates_subdir(client, tmp_path):
|
|
content = json.dumps({"checkpoint": "model.safetensors"})
|
|
resp = await client.post(
|
|
"/styles/save", json={"file": "anime/test.json", "content": content}
|
|
)
|
|
assert resp.status == 200
|
|
assert (tmp_path / "anime" / "test.json").exists()
|
|
|
|
|
|
# --- /styles/new ---
|
|
|
|
|
|
async def test_new_creates_file(client, tmp_path):
|
|
resp = await client.post("/styles/new", json={"file": "new_style"})
|
|
assert resp.status == 200
|
|
data = await resp.json()
|
|
assert data["ok"] is True
|
|
assert data["file"] == "new_style.json"
|
|
assert (tmp_path / "new_style.json").exists()
|
|
|
|
|
|
async def test_new_adds_json_extension(client):
|
|
resp = await client.post("/styles/new", json={"file": "mystyle"})
|
|
data = await resp.json()
|
|
assert data["file"].endswith(".json")
|
|
|
|
|
|
async def test_new_empty_name(client):
|
|
resp = await client.post("/styles/new", json={"file": ""})
|
|
assert resp.status == 400
|
|
|
|
|
|
async def test_new_conflict(client, tmp_path):
|
|
(tmp_path / "existing.json").write_text("{}")
|
|
resp = await client.post("/styles/new", json={"file": "existing.json"})
|
|
assert resp.status == 409
|
|
|
|
|
|
async def test_new_template_valid_json(client):
|
|
resp = await client.post("/styles/new", json={"file": "template_test"})
|
|
data = await resp.json()
|
|
parsed = json.loads(data["content"])
|
|
assert "checkpoint" in parsed
|
|
assert "loras" in parsed
|
|
assert isinstance(parsed["loras"], list)
|
|
|
|
|
|
async def test_new_invalid_path(client):
|
|
resp = await client.post("/styles/new", json={"file": "../../evil"})
|
|
assert resp.status == 400
|
|
|
|
|
|
# --- /styles/load ---
|
|
|
|
|
|
async def test_load_existing(client, tmp_path):
|
|
(tmp_path / "style.json").write_text('{"checkpoint": "model.safetensors"}')
|
|
resp = await client.get("/styles/load?file=style.json")
|
|
assert resp.status == 200
|
|
data = await resp.json()
|
|
assert "checkpoint" in data["content"]
|
|
|
|
|
|
async def test_load_not_found(client):
|
|
resp = await client.get("/styles/load?file=missing.json")
|
|
assert resp.status == 404
|
|
|
|
|
|
async def test_load_invalid_path(client):
|
|
resp = await client.get("/styles/load?file=../../etc/passwd")
|
|
assert resp.status == 400
|
|
|
|
|
|
# --- /styles/delete ---
|
|
|
|
|
|
async def test_delete_existing(client, tmp_path):
|
|
(tmp_path / "to_delete.json").write_text("{}")
|
|
resp = await client.post("/styles/delete", json={"file": "to_delete.json"})
|
|
assert resp.status == 200
|
|
data = await resp.json()
|
|
assert data["ok"] is True
|
|
assert not (tmp_path / "to_delete.json").exists()
|
|
|
|
|
|
async def test_delete_not_found(client):
|
|
resp = await client.post("/styles/delete", json={"file": "missing.json"})
|
|
assert resp.status == 404
|
|
|
|
|
|
async def test_delete_invalid_path(client):
|
|
resp = await client.post("/styles/delete", json={"file": "../../evil.json"})
|
|
assert resp.status == 400
|
|
|
|
|
|
# --- /styles/preview ---
|
|
|
|
|
|
async def test_preview_png_found(client, tmp_path):
|
|
# Create a minimal 1x1 PNG
|
|
png_bytes = (
|
|
b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x01"
|
|
b"\x00\x00\x00\x01\x08\x02\x00\x00\x00\x90wS\xde\x00\x00"
|
|
b"\x00\x0cIDATx\x9cc\xf8\x0f\x00\x00\x01\x01\x00\x05\x18"
|
|
b"\xd8N\x00\x00\x00\x00IEND\xaeB`\x82"
|
|
)
|
|
(tmp_path / "style.png").write_bytes(png_bytes)
|
|
resp = await client.get("/styles/preview?file=style.json")
|
|
assert resp.status == 200
|
|
data = await resp.json()
|
|
assert data["image"] is not None
|
|
assert data["image"].startswith("data:image/png;base64,")
|
|
|
|
|
|
async def test_preview_jpg_found(client, tmp_path):
|
|
# Minimal JPEG bytes
|
|
jpg_bytes = bytes(
|
|
[
|
|
0xFF,
|
|
0xD8,
|
|
0xFF,
|
|
0xE0,
|
|
0x00,
|
|
0x10,
|
|
0x4A,
|
|
0x46,
|
|
0x49,
|
|
0x46,
|
|
0x00,
|
|
0x01,
|
|
0x01,
|
|
0x00,
|
|
0x00,
|
|
0x01,
|
|
0x00,
|
|
0x01,
|
|
0x00,
|
|
0x00,
|
|
0xFF,
|
|
0xD9,
|
|
]
|
|
)
|
|
(tmp_path / "style.jpg").write_bytes(jpg_bytes)
|
|
resp = await client.get("/styles/preview?file=style.json")
|
|
assert resp.status == 200
|
|
data = await resp.json()
|
|
assert data["image"].startswith("data:image/jpeg;base64,")
|
|
|
|
|
|
async def test_preview_no_image(client, tmp_path):
|
|
(tmp_path / "style.json").write_text("{}")
|
|
resp = await client.get("/styles/preview?file=style.json")
|
|
assert resp.status == 200
|
|
data = await resp.json()
|
|
assert data["image"] is None
|
|
|
|
|
|
async def test_preview_invalid_path(client):
|
|
resp = await client.get("/styles/preview?file=../../etc/passwd")
|
|
assert resp.status == 400
|
|
|
|
|
|
async def test_preview_subdir(client, tmp_path):
|
|
sub = tmp_path / "anime"
|
|
sub.mkdir()
|
|
png_bytes = (
|
|
b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x01"
|
|
b"\x00\x00\x00\x01\x08\x02\x00\x00\x00\x90wS\xde\x00\x00"
|
|
b"\x00\x0cIDATx\x9cc\xf8\x0f\x00\x00\x01\x01\x00\x05\x18"
|
|
b"\xd8N\x00\x00\x00\x00IEND\xaeB`\x82"
|
|
)
|
|
(sub / "illustrious.png").write_bytes(png_bytes)
|
|
resp = await client.get("/styles/preview?file=anime/illustrious.json")
|
|
assert resp.status == 200
|
|
data = await resp.json()
|
|
assert data["image"] is not None
|