111 lines
2.7 KiB
Python
111 lines
2.7 KiB
Python
from __future__ import annotations
|
|
|
|
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
|