from __future__ import annotations import asyncio import os import time from pathlib import Path import pytest 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}" 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) 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() 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()