MpakoImageAIbot/tests/test_web.py

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 insert_user(
    path: Path,
    *,
    user_id: int,
    username: str | None = "alice",
    credits: int = 3,
    is_admin: int = 0,
    is_blocked: int = 0,
) -> None:
    connection = sqlite3.connect(path)
    try:
        connection.execute(
            """
            INSERT INTO users
                (user_id, username, credits, last_daily_credit, created_at,
                 is_admin, is_blocked)
            VALUES (?, ?, ?, '2026-10-18', '2026-10-18T12:00:00+00:00', ?, ?)
            """,
            (user_id, username, credits, is_admin, is_blocked),
        )
        connection.commit()
    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


def test_user_management(tmp_path: Path) -> None:
    settings = make_settings(tmp_path)
    app = create_app(settings)
    with TestClient(app) as client:
        insert_user(settings.db_path, user_id=1, username="alice")
        insert_user(settings.db_path, user_id=2, username="bob", credits=0)

        page = client.get("/users")
        assert page.status_code == 200
        assert "@alice" in page.text
        assert "@bob" in page.text

        assert client.post("/users/1/credits", data={"credits": "9"}).status_code == 200
        assert client.post("/users/1/admin").status_code == 200
        assert client.post("/users/2/block").status_code == 200

        connection = sqlite3.connect(settings.db_path)
        try:
            first = connection.execute(
                "SELECT credits, is_admin, is_blocked FROM users WHERE user_id = 1"
            ).fetchone()
            second = connection.execute(
                "SELECT is_blocked FROM users WHERE user_id = 2"
            ).fetchone()
        finally:
            connection.close()
        assert first == (9, 1, 0)
        assert second == (1,)


def test_user_search(tmp_path: Path) -> None:
    settings = make_settings(tmp_path)
    app = create_app(settings)
    with TestClient(app) as client:
        insert_user(settings.db_path, user_id=1, username="alice")
        insert_user(settings.db_path, user_id=2, username="bob")

        page = client.get("/users?q=alice")
        assert page.status_code == 200
        assert "@alice" in page.text
        assert "@bob" not in page.text