385 lines
11 KiB
Python
385 lines
11 KiB
Python
from __future__ import annotations
|
|
|
|
from datetime import timedelta
|
|
|
|
import pytest
|
|
|
|
from db import Database, UserAlreadyQueued, msk_today
|
|
|
|
|
|
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
|
|
|
|
assert await db.set_age_confirmed(1, True) is True
|
|
user = await db.get_user(1)
|
|
assert user is not None
|
|
assert user.age_confirmed is True
|
|
assert user.age_declined is False
|
|
assert await db.set_age_declined(1) is True
|
|
user = await db.get_user(1)
|
|
assert user is not None
|
|
assert user.age_confirmed is False
|
|
assert user.age_declined is True
|
|
|
|
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 = msk_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_hidden_requests_are_excluded_from_public_queries(
|
|
db: Database,
|
|
) -> None:
|
|
visible = await db.create_request(
|
|
user_id=1,
|
|
username="alice",
|
|
mode="t2i",
|
|
prompt_text="visible",
|
|
input_file_path=None,
|
|
seed=1,
|
|
)
|
|
await db.set_done(visible.id, "/tmp/visible.png")
|
|
hidden = await db.create_request(
|
|
user_id=2,
|
|
username="bob",
|
|
mode="t2i",
|
|
prompt_text="hidden",
|
|
input_file_path=None,
|
|
seed=2,
|
|
)
|
|
await db.set_done(hidden.id, "/tmp/hidden.png")
|
|
assert await db.set_hidden(hidden.id, True) is True
|
|
|
|
public = await db.list_requests(include_hidden=False)
|
|
assert [item.id for item in public] == [visible.id]
|
|
assert await db.count_requests(include_hidden=False) == 1
|
|
|
|
admin = await db.list_requests(include_hidden=True)
|
|
assert {item.id for item in admin} == {visible.id, hidden.id}
|
|
|
|
stats = await db.stats(include_hidden=False)
|
|
assert stats["total"] == 1
|
|
assert stats["users"] == 1
|
|
|
|
public_users = await db.list_users(include_hidden=False)
|
|
assert {item.user_id for item in public_users} == {1}
|
|
|
|
loaded = await db.get_request(hidden.id)
|
|
assert loaded is not None
|
|
assert loaded.hidden is True
|
|
|
|
|
|
async def test_generation_quality_setting(db: Database) -> None:
|
|
user = await db.ensure_user(1, "alice")
|
|
assert user.generation_quality == "standard"
|
|
assert user.generation_steps == 12
|
|
|
|
assert await db.set_generation_quality(1, "improved") is True
|
|
assert await db.set_generation_steps(1, 8) is True
|
|
user = await db.get_user(1)
|
|
assert user is not None
|
|
assert user.generation_quality == "improved"
|
|
assert user.generation_steps == 8
|
|
|
|
with pytest.raises(ValueError):
|
|
await db.set_generation_quality(1, "unknown")
|
|
with pytest.raises(ValueError):
|
|
await db.set_generation_steps(1, 9)
|
|
|
|
|
|
async def test_user_context_lifecycle(db: Database) -> None:
|
|
assert await db.get_user_context(1) is None
|
|
|
|
await db.save_user_context(
|
|
user_id=1,
|
|
base_prompt="A red apple",
|
|
current_prompt="A red apple on a table",
|
|
last_request_id=10,
|
|
last_mode="t2i",
|
|
)
|
|
context = await db.get_user_context(1)
|
|
assert context is not None
|
|
assert context.base_prompt == "A red apple"
|
|
assert context.current_prompt == "A red apple on a table"
|
|
assert context.last_request_id == 10
|
|
assert context.last_mode == "t2i"
|
|
|
|
await db.save_user_context(
|
|
user_id=1,
|
|
base_prompt="A red apple",
|
|
current_prompt="A red apple on a blue table",
|
|
last_request_id=11,
|
|
last_mode="t2i",
|
|
)
|
|
context = await db.get_user_context(1)
|
|
assert context is not None
|
|
assert context.base_prompt == "A red apple"
|
|
assert context.current_prompt == "A red apple on a blue table"
|
|
assert context.last_request_id == 11
|
|
|
|
assert await db.clear_user_context(1) is True
|
|
assert await db.get_user_context(1) is None
|
|
assert await db.clear_user_context(1) is False
|
|
|
|
|
|
async def test_ensure_user_updates_username_without_resetting_credits(
|
|
db: Database,
|
|
) -> None:
|
|
user = await db.ensure_user(1, "alice")
|
|
assert user.credits == 3
|
|
|
|
assert await db.add_credits(1, 2) == 5
|
|
user = await db.ensure_user(1, "alice_new")
|
|
assert user.username == "alice_new"
|
|
assert user.credits == 5
|
|
|
|
user = await db.ensure_user(1)
|
|
assert user.username == "alice_new"
|
|
assert user.credits == 5
|