MpakoImageAIbot/tests/test_web.py

from __future__ import annotations

from datetime import UTC, date, datetime

import psycopg
from fastapi.testclient import TestClient
from PIL import Image
from psycopg.rows import dict_row

from config import Settings
from web import create_app


def insert_request(
    database_url: str,
    *,
    user_id: int,
    prompt: str,
    status: str = "done",
    mode: str = "t2i",
    output: str | None = None,
    username: str | None = None,
    enhanced: str | None = None,
) -> int:
    with psycopg.connect(database_url, row_factory=dict_row) as connection:
        row = connection.execute(
            """
            INSERT INTO requests
                (user_id, username, timestamp, prompt_text, enhanced_prompt,
                 status, mode, seed, output_file_path)
            VALUES (%s, %s, %s, %s, %s, %s, %s, 1, %s)
            RETURNING id
            """,
            (
                user_id,
                username,
                datetime(2026, 10, 18, 12, tzinfo=UTC),
                prompt,
                enhanced,
                status,
                mode,
                output,
            ),
        ).fetchone()
    if row is None:
        raise RuntimeError("Не удалось вставить запрос")
    return int(row["id"])


def insert_user(
    database_url: str,
    *,
    user_id: int,
    username: str | None = "alice",
    credits: int = 3,
    is_admin: bool = False,
    is_blocked: bool = False,
) -> None:
    with psycopg.connect(database_url) as connection:
        connection.execute(
            """
            INSERT INTO users
                (user_id, username, credits, last_daily_credit, created_at,
                 is_admin, is_blocked)
            VALUES (%s, %s, %s, %s, %s, %s, %s)
            ON CONFLICT (user_id) DO NOTHING
            """,
            (
                user_id,
                username,
                credits,
                date(2026, 10, 18),
                datetime(2026, 10, 18, 12, tzinfo=UTC),
                is_admin,
                is_blocked,
            ),
        )


def test_history_page_and_user_filter(web_settings: Settings) -> None:
    app = create_app(web_settings)
    with TestClient(app) as client:
        insert_request(
            web_settings.database_url,
            user_id=1,
            username="alice",
            prompt="apple prompt",
        )
        insert_request(
            web_settings.database_url,
            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(web_settings: Settings) -> None:
    output = web_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(web_settings)
    with TestClient(app) as client:
        request_id = insert_request(
            web_settings.database_url,
            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(web_settings: Settings) -> None:
    app = create_app(web_settings)
    with TestClient(app) as client:
        insert_user(
            web_settings.database_url, user_id=1, username="alice"
        )
        insert_user(
            web_settings.database_url,
            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

        with psycopg.connect(
            web_settings.database_url, row_factory=dict_row
        ) as connection:
            first = connection.execute(
                "SELECT credits, is_admin, is_blocked FROM users"
                " WHERE user_id = %s",
                (1,),
            ).fetchone()
            second = connection.execute(
                "SELECT is_blocked FROM users WHERE user_id = %s", (2,)
            ).fetchone()
        assert first is not None
        assert second is not None
        assert first["credits"] == 9
        assert first["is_admin"] is True
        assert first["is_blocked"] is False
        assert second["is_blocked"] is True


def test_user_search(web_settings: Settings) -> None:
    app = create_app(web_settings)
    with TestClient(app) as client:
        insert_user(
            web_settings.database_url, user_id=1, username="alice"
        )
        insert_user(
            web_settings.database_url, 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