MpakoImageAIbot/tests/test_web.py

115 lines
3.6 KiB
Python

from __future__ import annotations
import sqlite3
from pathlib import Path
from fastapi.testclient import TestClient
from PIL import Image
from config import Settings
from web import create_app
def make_settings(tmp_path: Path, **overrides) -> Settings:
values = {
"bot_token": "",
"db_path": tmp_path / "bot.db",
"storage_dir": tmp_path / "storage",
}
values.update(overrides)
return Settings(**values)
def insert_request(
path: Path,
*,
user_id: int,
prompt: str,
status: str = "done",
mode: str = "t2i",
output: str | None = None,
username: str | None = None,
enhanced: str | None = None,
) -> int:
connection = sqlite3.connect(path)
try:
cursor = connection.execute(
"""
INSERT INTO requests
(user_id, username, timestamp, prompt_text, enhanced_prompt,
status, mode, seed, output_file_path)
VALUES (?, ?, '2026-10-18T12:00:00+00:00', ?, ?, ?, ?, 1, ?)
""",
(user_id, username, prompt, enhanced, status, mode, output),
)
connection.commit()
return int(cursor.lastrowid)
finally:
connection.close()
def test_history_page_and_user_filter(tmp_path: Path) -> None:
settings = make_settings(tmp_path)
app = create_app(settings)
with TestClient(app) as client:
insert_request(
settings.db_path, user_id=1, username="alice", prompt="apple prompt"
)
insert_request(settings.db_path, user_id=2, prompt="banana prompt")
response = client.get("/")
assert response.status_code == 200
assert "apple prompt" in response.text
assert "banana prompt" in response.text
assert "@alice" in response.text
assert 'title="ID: 2">2</a>' in response.text
response = client.get("/?user_id=1")
assert response.status_code == 200
assert "apple prompt" in response.text
assert "banana prompt" not in response.text
response = client.get("/?user_id=abc")
assert response.status_code == 200
assert "user_id должен состоять только из цифр" in response.text
def test_detail_and_media(tmp_path: Path) -> None:
settings = make_settings(tmp_path)
output = settings.storage_dir / "1" / "1_output.png"
output.parent.mkdir(parents=True)
with Image.new("RGB", (320, 200), (180, 40, 40)) as image:
image.save(output)
app = create_app(settings)
with TestClient(app) as client:
request_id = insert_request(
settings.db_path,
user_id=1,
prompt="detail prompt",
enhanced="An enhanced English prompt",
output=str(output),
)
detail = client.get(f"/request/{request_id}")
assert detail.status_code == 200
assert "detail prompt" in detail.text
assert "An enhanced English prompt" in detail.text
index = client.get("/")
assert f"/thumb/{request_id}" in index.text
assert "detail prompt" in index.text
assert "An enhanced English prompt" in index.text
media = client.get(f"/media/{request_id}/output")
assert media.status_code == 200
assert media.content.startswith(b"\x89PNG")
thumb = client.get(f"/thumb/{request_id}")
assert thumb.status_code == 200
assert thumb.headers["content-type"] == "image/jpeg"
assert thumb.content.startswith(b"\xff\xd8")
missing = client.get(f"/media/{request_id}/input")
assert missing.status_code == 404