MpakoImageAIbot/tests/test_db.py

from __future__ import annotations

import sqlite3

import pytest
import pytest_asyncio

from db import Database, UserAlreadyQueued


@pytest_asyncio.fixture
async def db(tmp_path):
    database = Database(tmp_path / "bot.db")
    await database.connect()
    yield database
    await database.close()


async def test_create_and_read(db: Database) -> None:
    request = await db.create_request(
        user_id=1,
        mode="t2i",
        prompt_text="apple",
        input_file_path=None,
        seed=42,
    )
    assert request.id > 0
    assert request.status == "queued"
    assert request.mode == "t2i"
    assert request.seed == 42

    loaded = await db.get_request(request.id)
    assert loaded == request


async def test_one_active_request_per_user(db: Database) -> None:
    first = await db.create_request(
        user_id=1,
        mode="t2i",
        prompt_text="first",
        input_file_path=None,
        seed=1,
    )
    with pytest.raises(UserAlreadyQueued):
        await db.create_request(
            user_id=1,
            mode="t2i",
            prompt_text="second",
            input_file_path=None,
            seed=2,
        )

    await db.set_done(first.id, "/tmp/out.png")
    second = await db.create_request(
        user_id=1,
        mode="t2i",
        prompt_text="second",
        input_file_path=None,
        seed=3,
    )
    assert second.id != first.id


async def test_positions(db: Database) -> None:
    first = await db.create_request(
        user_id=1,
        mode="t2i",
        prompt_text="first",
        input_file_path=None,
        seed=1,
    )
    second = await db.create_request(
        user_id=2,
        mode="t2i",
        prompt_text="second",
        input_file_path=None,
        seed=2,
    )
    assert await db.count_ahead(first.id) == 0
    assert await db.count_ahead(second.id) == 1

    await db.set_running(first.id, "prompt-1")
    assert await db.count_ahead(second.id) == 1

    await db.set_done(first.id, "/tmp/out.png")
    assert await db.count_ahead(second.id) == 0


async def test_recover_marks_running_failed_and_returns_queued(db: Database) -> None:
    running = await db.create_request(
        user_id=1,
        mode="t2i",
        prompt_text="running",
        input_file_path=None,
        seed=1,
    )
    queued = await db.create_request(
        user_id=2,
        mode="t2i",
        prompt_text="queued",
        input_file_path=None,
        seed=2,
    )
    await db.set_running(running.id, "prompt-1")

    recovered = await db.recover()

    assert [item.id for item in recovered] == [queued.id]
    failed = await db.get_request(running.id)
    assert failed is not None
    assert failed.status == "failed"
    assert failed.finished_at is not None


async def test_history_filters_and_stats(db: Database) -> None:
    first = await db.create_request(
        user_id=1,
        username="alice",
        mode="t2i",
        prompt_text="apple",
        input_file_path=None,
        seed=1,
    )
    await db.set_done(first.id, "/tmp/1.png")
    second = await db.create_request(
        user_id=1,
        username="alice",
        mode="i2i",
        prompt_text="banana",
        input_file_path="/tmp/in.png",
        seed=2,
    )
    await db.set_failed(second.id, "error")
    third = await db.create_request(
        user_id=2,
        mode="t2i",
        prompt_text="cherry",
        input_file_path=None,
        seed=3,
    )

    by_user = await db.list_requests(user_id=1)
    assert {item.id for item in by_user} == {first.id, second.id}

    failed = await db.list_requests(status="failed")
    assert [item.id for item in failed] == [second.id]

    queued_t2i = await db.list_requests(mode="t2i", status="queued")
    assert [item.id for item in queued_t2i] == [third.id]

    assert await db.count_requests(user_id=1) == 2
    assert await db.count_requests(status="failed") == 1

    stats = await db.stats()
    assert stats["total"] == 3
    assert stats["users"] == 2
    assert stats["by_status"]["done"] == 1
    assert stats["by_status"]["failed"] == 1
    assert stats["by_status"]["queued"] == 1

    users = await db.list_users()
    assert {user.user_id for user in users} == {1, 2}
    alice = next(user for user in users if user.user_id == 1)
    assert alice.username == "alice"
    assert alice.display_user == "@alice"


async def test_migrates_legacy_schema_without_username(tmp_path) -> None:
    path = tmp_path / "legacy.db"
    connection = sqlite3.connect(path)
    connection.execute(
        """
        CREATE TABLE requests (
            id               INTEGER PRIMARY KEY AUTOINCREMENT,
            user_id          INTEGER NOT NULL,
            timestamp        TEXT    NOT NULL,
            prompt_text      TEXT    NOT NULL,
            input_file_path  TEXT,
            output_file_path TEXT,
            status           TEXT    NOT NULL DEFAULT 'queued'
                             CHECK (status IN ('queued', 'running', 'done', 'failed')),
            comfy_prompt_id  TEXT,
            mode             TEXT    NOT NULL CHECK (mode IN ('t2i', 'i2i')),
            seed             INTEGER,
            error_text       TEXT,
            finished_at      TEXT
        )
        """
    )
    connection.commit()
    connection.close()

    database = Database(path)
    await database.connect()
    try:
        request = await database.create_request(
            user_id=1,
            username="alice",
            mode="t2i",
            prompt_text="apple",
            input_file_path=None,
            seed=1,
        )
        assert request.username == "alice"
        await database.set_enhanced_prompt(request.id, "a red apple")
        loaded = await database.get_request(request.id)
        assert loaded is not None
        assert loaded.enhanced_prompt == "a red apple"
    finally:
        await database.close()