254 lines
7.4 KiB
Python
254 lines
7.4 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import time
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
import pytest_asyncio
|
|
|
|
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]] = []
|
|
|
|
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
|
|
|
|
|
|
class FakeEnhancer:
|
|
enabled = True
|
|
|
|
def __init__(self) -> None:
|
|
self.calls: list[tuple[str, str]] = []
|
|
|
|
async def enhance(self, prompt: str, mode: str) -> str:
|
|
self.calls.append((prompt, mode))
|
|
return f"enhanced: {prompt}"
|
|
|
|
|
|
@pytest_asyncio.fixture
|
|
async def db(tmp_path):
|
|
database = Database(tmp_path / "bot.db")
|
|
await database.connect()
|
|
yield database
|
|
await database.close()
|
|
|
|
|
|
def make_settings(tmp_path: Path, **overrides) -> Settings:
|
|
values = {
|
|
"bot_token": "test",
|
|
"storage_dir": tmp_path / "storage",
|
|
"db_path": tmp_path / "bot.db",
|
|
"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)
|
|
|
|
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"
|
|
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"
|
|
input_path.write_bytes(b"input")
|
|
|
|
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"
|
|
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_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()
|