318 lines
9.1 KiB
Python
318 lines
9.1 KiB
Python
from __future__ import annotations
|
|
|
|
import sqlite3
|
|
from datetime import date, timedelta
|
|
|
|
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_premium_queue_priority_counts(db: Database) -> None:
|
|
normal = await db.create_request(
|
|
user_id=1,
|
|
mode="t2i",
|
|
prompt_text="normal",
|
|
input_file_path=None,
|
|
seed=1,
|
|
)
|
|
premium = await db.create_request(
|
|
user_id=2,
|
|
mode="t2i",
|
|
prompt_text="premium",
|
|
input_file_path=None,
|
|
seed=2,
|
|
is_premium=True,
|
|
)
|
|
counts = await db.queue_counts()
|
|
assert counts == {"premium": 1, "normal": 1}
|
|
assert await db.count_ahead(premium.id) == 0
|
|
assert await db.count_ahead(normal.id) == 1
|
|
|
|
await db.set_done(normal.id, "/tmp/normal.png")
|
|
counts = await db.queue_counts()
|
|
assert counts == {"premium": 1, "normal": 0}
|
|
|
|
|
|
async def test_pending_payment_lifecycle(db: Database) -> None:
|
|
pending = await db.create_pending_payment(
|
|
payload="payload-1",
|
|
user_id=1,
|
|
mode="t2i",
|
|
prompt="apple",
|
|
username="alice",
|
|
input_file_path=None,
|
|
amount=0,
|
|
)
|
|
assert pending.status == "pending"
|
|
assert pending.amount == 0
|
|
|
|
assert await db.set_pending_choice("payload-1", True, 15) is True
|
|
loaded = await db.get_pending_payment("payload-1")
|
|
assert loaded is not None
|
|
assert loaded.is_premium is True
|
|
assert loaded.amount == 15
|
|
|
|
assert await db.mark_pending_paid("payload-1", "charge-1") is True
|
|
loaded = await db.get_pending_payment("payload-1")
|
|
assert loaded is not None
|
|
assert loaded.status == "paid"
|
|
assert loaded.telegram_payment_charge_id == "charge-1"
|
|
|
|
await db.attach_pending_request("payload-1", 42)
|
|
loaded = await db.get_pending_payment("payload-1")
|
|
assert loaded is not None
|
|
assert loaded.request_id == 42
|
|
|
|
await db.mark_pending_refunded("payload-1")
|
|
loaded = await db.get_pending_payment("payload-1")
|
|
assert loaded is not None
|
|
assert loaded.status == "refunded"
|
|
|
|
request = await db.create_request(
|
|
user_id=2,
|
|
mode="t2i",
|
|
prompt_text="paid apple",
|
|
input_file_path=None,
|
|
seed=1,
|
|
payment_charge_id="charge-2",
|
|
)
|
|
assert request.payment_charge_id == "charge-2"
|
|
assert request.payment_refunded is False
|
|
await db.set_refunded(request.id)
|
|
loaded_request = await db.get_request(request.id)
|
|
assert loaded_request is not None
|
|
assert loaded_request.payment_refunded is True
|
|
|
|
|
|
async def test_user_credits_referral_and_daily_accrual(db: Database) -> None:
|
|
user = await db.ensure_user(1, "alice")
|
|
assert user.credits == 3
|
|
assert user.last_daily_credit
|
|
|
|
assert await db.add_credits(1, 2) == 5
|
|
assert await db.consume_credit(1) == 4
|
|
assert await db.consume_credit(1) == 3
|
|
|
|
friend = await db.ensure_user(2, "bob")
|
|
assert friend.credits == 3
|
|
assert await db.set_referrer(2, 1) is True
|
|
assert await db.set_referrer(2, 1) is False
|
|
assert await db.mark_first_payment(2) is True
|
|
assert await db.mark_first_payment(2) is False
|
|
|
|
friend = await db.get_user(2)
|
|
assert friend is not None
|
|
assert friend.referrer_id == 1
|
|
assert friend.first_paid_at is not None
|
|
|
|
future = date.today() + timedelta(days=2)
|
|
assert await db.accrue_daily_credits(future) >= 2
|
|
user = await db.get_user(1)
|
|
assert user is not None
|
|
assert user.credits == 5
|
|
assert await db.accrue_daily_credits(future) == 0
|
|
|
|
|
|
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()
|