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()