348 lines
10 KiB
Python
348 lines
10 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import os
|
|
import time
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
from PIL import Image
|
|
|
|
from config import Settings
|
|
from db import Database, UserAlreadyQueued
|
|
from moderation import PromptBlocked
|
|
from queue_manager import QueueEvent, QueueFull, QueueManager
|
|
from workflows import WorkflowBuilder
|
|
|
|
BASE = Path(__file__).resolve().parents[1] / "workflows"
|
|
|
|
|
|
class FakeComfy:
|
|
def __init__(self) -> None:
|
|
self.workflows: list[dict] = []
|
|
self.uploads: list[tuple[Path, str]] = []
|
|
self.deleted_history: list[str] = []
|
|
|
|
async def open_progress(self, client_id: str):
|
|
return None
|
|
|
|
async def submit(self, workflow: dict, client_id: str) -> str:
|
|
self.workflows.append(workflow)
|
|
return f"prompt-{len(self.workflows)}"
|
|
|
|
async def wait_for_history(self, prompt_id: str, *, timeout: float) -> dict:
|
|
return {"outputs": {"9": {"images": [{"filename": "x.png", "type": "output"}]}}}
|
|
|
|
def find_output_image(self, entry: dict) -> dict | None:
|
|
images = entry["outputs"]["9"]["images"]
|
|
return images[0] if images else None
|
|
|
|
async def download_output(self, entry: dict, destination: Path) -> Path:
|
|
destination.parent.mkdir(parents=True, exist_ok=True)
|
|
destination.write_bytes(b"png")
|
|
return destination
|
|
|
|
async def upload_image(self, path: Path, remote_name: str) -> str:
|
|
self.uploads.append((path, remote_name))
|
|
return remote_name
|
|
|
|
async def delete_history(self, prompt_id: str) -> None:
|
|
self.deleted_history.append(prompt_id)
|
|
|
|
|
|
class FakeEnhancer:
|
|
enabled = True
|
|
|
|
def __init__(self) -> None:
|
|
self.calls: list[tuple[str, str]] = []
|
|
self.contexts: list[tuple[str | None, str | None]] = []
|
|
|
|
async def enhance(
|
|
self,
|
|
prompt: str,
|
|
mode: str,
|
|
*,
|
|
previous_prompt: str | None = None,
|
|
base_prompt: str | None = None,
|
|
) -> str:
|
|
self.calls.append((prompt, mode))
|
|
self.contexts.append((previous_prompt, base_prompt))
|
|
return f"enhanced: {prompt}"
|
|
|
|
|
|
def make_settings(tmp_path: Path, **overrides) -> Settings:
|
|
values = {
|
|
"bot_token": "test",
|
|
"storage_dir": tmp_path / "storage",
|
|
"database_url": os.getenv("TEST_DATABASE_URL", ""),
|
|
"workflow_t2i": BASE / "txt2img.json",
|
|
"workflow_i2i": BASE / "img2img.json",
|
|
"workflow_mapping": BASE / "mapping.json",
|
|
}
|
|
values.update(overrides)
|
|
return Settings(**values)
|
|
|
|
|
|
def make_queue(
|
|
db: Database,
|
|
settings: Settings,
|
|
events: list[QueueEvent],
|
|
*,
|
|
enhancer=None,
|
|
) -> QueueManager:
|
|
async def on_event(event: QueueEvent) -> None:
|
|
events.append(event)
|
|
|
|
return QueueManager(
|
|
db=db,
|
|
comfy=FakeComfy(),
|
|
workflows=WorkflowBuilder(
|
|
settings.workflow_t2i,
|
|
settings.workflow_i2i,
|
|
settings.workflow_mapping,
|
|
),
|
|
settings=settings,
|
|
on_event=on_event,
|
|
enhancer=enhancer,
|
|
)
|
|
|
|
|
|
async def wait_for_event(
|
|
events: list[QueueEvent], kind: str, timeout: float = 5.0
|
|
) -> QueueEvent:
|
|
deadline = time.monotonic() + timeout
|
|
while time.monotonic() < deadline:
|
|
for event in events:
|
|
if event.kind == kind:
|
|
return event
|
|
await asyncio.sleep(0.01)
|
|
raise AssertionError(f"Событие {kind!r} не пришло: {[e.kind for e in events]}")
|
|
|
|
|
|
async def test_positions_and_duplicate_user(db: Database, tmp_path: Path) -> None:
|
|
events: list[QueueEvent] = []
|
|
queue = make_queue(db, make_settings(tmp_path), events)
|
|
|
|
first, first_position = await queue.enqueue(
|
|
user_id=1, mode="t2i", prompt="first"
|
|
)
|
|
second, second_position = await queue.enqueue(
|
|
user_id=2, mode="t2i", prompt="second"
|
|
)
|
|
|
|
assert first_position == 0
|
|
assert second_position == 1
|
|
with pytest.raises(UserAlreadyQueued):
|
|
await queue.enqueue(user_id=1, mode="t2i", prompt="again")
|
|
|
|
await queue.stop()
|
|
assert first.id != second.id
|
|
|
|
|
|
async def test_queue_limit(db: Database, tmp_path: Path) -> None:
|
|
events: list[QueueEvent] = []
|
|
queue = make_queue(
|
|
db,
|
|
make_settings(tmp_path, max_queue_size=1),
|
|
events,
|
|
)
|
|
|
|
await queue.enqueue(user_id=1, mode="t2i", prompt="first")
|
|
with pytest.raises(QueueFull):
|
|
await queue.enqueue(user_id=2, mode="t2i", prompt="second")
|
|
|
|
await queue.stop()
|
|
|
|
|
|
async def test_text_to_image_job(db: Database, tmp_path: Path) -> None:
|
|
events: list[QueueEvent] = []
|
|
settings = make_settings(tmp_path)
|
|
queue = make_queue(db, settings, events)
|
|
|
|
comfy = queue._comfy # type: ignore[attr-defined]
|
|
await queue.start()
|
|
try:
|
|
request, _ = await queue.enqueue(
|
|
user_id=1, mode="t2i", prompt="apple"
|
|
)
|
|
event = await wait_for_event(events, "done")
|
|
|
|
assert event.request.id == request.id
|
|
assert event.request.status == "done"
|
|
output = Path(event.request.output_file_path or "")
|
|
assert output.exists()
|
|
assert output.name == f"{request.id}_output.png"
|
|
assert comfy.deleted_history == ["prompt-1"]
|
|
finally:
|
|
await queue.stop()
|
|
|
|
|
|
async def test_image_to_image_uploads_input(db: Database, tmp_path: Path) -> None:
|
|
events: list[QueueEvent] = []
|
|
settings = make_settings(tmp_path)
|
|
queue = make_queue(db, settings, events)
|
|
comfy = queue._comfy # type: ignore[attr-defined]
|
|
|
|
input_path = tmp_path / "input.png"
|
|
Image.new("RGB", (64, 64), (120, 40, 40)).save(input_path)
|
|
await db.ensure_user(1, "alice")
|
|
await db.set_generation_quality(1, "improved")
|
|
|
|
await queue.start()
|
|
try:
|
|
request, _ = await queue.enqueue(
|
|
user_id=1,
|
|
mode="i2i",
|
|
prompt="make blue",
|
|
input_file_path=str(input_path),
|
|
)
|
|
await wait_for_event(events, "done")
|
|
|
|
assert comfy.uploads
|
|
assert comfy.uploads[0][0] == input_path
|
|
assert comfy.uploads[0][1] == f"{request.id}_input.png"
|
|
source = comfy.workflows[0]["20"]["inputs"]["image"]
|
|
assert source == f"{request.id}_input.png"
|
|
assert comfy.workflows[0]["25"]["inputs"]["vae"] == ["19", 2]
|
|
assert (
|
|
comfy.workflows[0]["19"]["inputs"]["unet_name"]
|
|
== settings.krea2_q6_model
|
|
)
|
|
assert comfy.workflows[0]["15"]["inputs"]["width"] == 1024
|
|
assert comfy.workflows[0]["15"]["inputs"]["height"] == 1024
|
|
finally:
|
|
await queue.stop()
|
|
|
|
|
|
async def test_prompt_enhancer_is_used(db: Database, tmp_path: Path) -> None:
|
|
events: list[QueueEvent] = []
|
|
settings = make_settings(tmp_path)
|
|
enhancer = FakeEnhancer()
|
|
queue = make_queue(db, settings, events, enhancer=enhancer)
|
|
comfy = queue._comfy # type: ignore[attr-defined]
|
|
|
|
await queue.start()
|
|
try:
|
|
request, _ = await queue.enqueue(
|
|
user_id=1, mode="t2i", prompt="apple"
|
|
)
|
|
await wait_for_event(events, "done")
|
|
|
|
assert enhancer.calls == [("apple", "t2i")]
|
|
assert comfy.workflows[0]["67"]["inputs"]["text"] == "enhanced: apple"
|
|
stored = await db.get_request(request.id)
|
|
assert stored is not None
|
|
assert stored.enhanced_prompt == "enhanced: apple"
|
|
assert any(event.kind == "enhancing" for event in events)
|
|
finally:
|
|
await queue.stop()
|
|
|
|
|
|
async def test_context_is_passed_to_enhancer(
|
|
db: Database, tmp_path: Path
|
|
) -> None:
|
|
events: list[QueueEvent] = []
|
|
enhancer = FakeEnhancer()
|
|
queue = make_queue(
|
|
db, make_settings(tmp_path), events, enhancer=enhancer
|
|
)
|
|
|
|
await queue.start()
|
|
try:
|
|
request, _ = await queue.enqueue(
|
|
user_id=1,
|
|
mode="t2i",
|
|
prompt="make it blue",
|
|
context_base_prompt="A red apple on a table",
|
|
context_prompt="A red apple on a table, soft light",
|
|
)
|
|
await wait_for_event(events, "done")
|
|
|
|
assert enhancer.contexts == [
|
|
(
|
|
"A red apple on a table, soft light",
|
|
"A red apple on a table",
|
|
)
|
|
]
|
|
stored = await db.get_request(request.id)
|
|
assert stored is not None
|
|
assert stored.context_prompt == "A red apple on a table, soft light"
|
|
assert stored.context_base_prompt == "A red apple on a table"
|
|
finally:
|
|
await queue.stop()
|
|
|
|
|
|
async def test_blocked_prompt_is_logged(db: Database, tmp_path: Path) -> None:
|
|
events: list[QueueEvent] = []
|
|
queue = make_queue(db, make_settings(tmp_path), events)
|
|
|
|
with pytest.raises(PromptBlocked):
|
|
await queue.enqueue(user_id=1, mode="t2i", prompt="child nude")
|
|
|
|
blocked = await db.list_blocked()
|
|
assert len(blocked) == 1
|
|
assert blocked[0].user_id == 1
|
|
assert blocked[0].reason == "csam"
|
|
assert blocked[0].prompt == "child nude"
|
|
|
|
await queue.stop()
|
|
|
|
|
|
async def test_admin_prompt_is_not_filtered(db: Database, tmp_path: Path) -> None:
|
|
events: list[QueueEvent] = []
|
|
queue = make_queue(db, make_settings(tmp_path), events)
|
|
|
|
request, _ = await queue.enqueue(
|
|
user_id=1,
|
|
mode="t2i",
|
|
prompt="child nude",
|
|
is_admin=True,
|
|
)
|
|
assert request.id > 0
|
|
assert await db.list_blocked() == []
|
|
|
|
await queue.stop()
|
|
|
|
|
|
async def test_premium_queue_goes_first(db: Database, tmp_path: Path) -> None:
|
|
events: list[QueueEvent] = []
|
|
queue = make_queue(db, make_settings(tmp_path), events)
|
|
|
|
normal, _ = await queue.enqueue(
|
|
user_id=1, mode="t2i", prompt="normal"
|
|
)
|
|
premium, _ = await queue.enqueue(
|
|
user_id=2, mode="t2i", prompt="premium", is_premium=True
|
|
)
|
|
|
|
first = await queue._next_request() # type: ignore[attr-defined]
|
|
second = await queue._next_request() # type: ignore[attr-defined]
|
|
assert first == premium.id
|
|
assert second == normal.id
|
|
|
|
await queue.stop()
|
|
|
|
|
|
async def test_premium_and_normal_alternate(db: Database, tmp_path: Path) -> None:
|
|
events: list[QueueEvent] = []
|
|
queue = make_queue(db, make_settings(tmp_path), events)
|
|
|
|
normal_1, _ = await queue.enqueue(user_id=1, mode="t2i", prompt="n1")
|
|
normal_2, _ = await queue.enqueue(user_id=2, mode="t2i", prompt="n2")
|
|
premium_1, _ = await queue.enqueue(
|
|
user_id=3, mode="t2i", prompt="p1", is_premium=True
|
|
)
|
|
premium_2, _ = await queue.enqueue(
|
|
user_id=4, mode="t2i", prompt="p2", is_premium=True
|
|
)
|
|
|
|
order = [await queue._next_request() for _ in range(4)] # type: ignore[attr-defined]
|
|
assert order == [
|
|
premium_1.id,
|
|
normal_1.id,
|
|
premium_2.id,
|
|
normal_2.id,
|
|
]
|
|
|
|
await queue.stop()
|