diff --git a/.env.example b/.env.example index 1ed16aa..d20fe19 100644 --- a/.env.example +++ b/.env.example @@ -43,6 +43,7 @@ BLOCKED_PUBLIC_FIGURES= # Оплата Telegram Stars PAYMENTS_ENABLED=true STARS_PRICE=5 +PREMIUM_STARS_PRICE=15 # Веб-история: только локальная сеть, без авторизации WEB_BIND_IP=192.168.31.158 diff --git a/README.md b/README.md index 6aa22ca..d610427 100644 --- a/README.md +++ b/README.md @@ -10,7 +10,8 @@ Telegram-бот для генерации и редактирования изо - **Веб-история**: просмотр запросов с фильтрами по `user_id`, статусу и режиму, просмотр изображений. - **Улучшение промптов**: перевод и переписывание запроса на английский через `TextGenerate` + `qwen3vl_4b` в ComfyUI. - **Фильтрация промптов**: блокировка CSAM, зоофилии, NSFW с публичными людьми и обходных вариантов для не-админов. -- **Оплата**: 5 Telegram Stars за генерацию, админам бесплатно, возврат при ошибке. +- **Оплата**: 5 Telegram Stars за обычную очередь, 15 за премиум; админам бесплатно, возврат при ошибке. +- **Премиум-очередь**: чередуется с обычной (одна премиум, одна обычная); если одна пуста — работает другая. - Глобальная очередь с одной активной задачей на пользователя. - Позиция в очереди и прогресс генерации через WebSocket ComfyUI. - SQLite: запросы, статусы, seed, пути к файлам. @@ -124,14 +125,18 @@ BLOCKED_PUBLIC_FIGURES= ## Оплата -Генерация стоит 5 Telegram Stars: +Генерация стоит 5 Telegram Stars в обычной очереди и 15 звёзд в премиум-очереди: ```dotenv PAYMENTS_ENABLED=true STARS_PRICE=5 +PREMIUM_STARS_PRICE=15 ``` -- Админы из `ADMIN_USER_IDS` генерируют бесплатно. +- Перед оплатой бот показывает количество людей в обычной и премиум-очередях. +- Пользователь выбирает очередь кнопкой, после чего получает счёт. +- Премиум- и обычные задачи чередуются: одна премиум, одна обычная; если одна очередь пуста, берётся из другой. +- Админы из `ADMIN_USER_IDS` генерируют бесплатно и попадают в премиум-очередь. - Не-админу выставляется счёт через Telegram Stars (`XTR`). - Фильтрация промпта выполняется до выставления счёта. - Если поставить задачу в очередь не удалось, звёзды возвращаются. diff --git a/bot.py b/bot.py index e16215f..57e3c27 100644 --- a/bot.py +++ b/bot.py @@ -6,11 +6,18 @@ import time import uuid from pathlib import Path -from telegram import InputFile, LabeledPrice, Update +from telegram import ( + InlineKeyboardButton, + InlineKeyboardMarkup, + InputFile, + LabeledPrice, + Update, +) from telegram.constants import ChatAction, ChatType from telegram.error import BadRequest, NetworkError, TelegramError, TimedOut from telegram.ext import ( Application, + CallbackQueryHandler, CommandHandler, ContextTypes, MessageHandler, @@ -71,6 +78,9 @@ class MpakoBot: application.add_handler( PreCheckoutQueryHandler(self._pre_checkout) ) + application.add_handler( + CallbackQueryHandler(self._queue_choice, pattern=r"^queue:") + ) application.add_handler( MessageHandler(filters.SUCCESSFUL_PAYMENT, self._successful_payment) ) @@ -120,8 +130,14 @@ class MpakoBot: def _start_text(self) -> str: if self._settings.payments_enabled: - price = self._settings.stars_price - return f"Стоимость одной генерации: {price} звёзд.\n\n" + START_TEXT + normal = self._settings.stars_price + premium = self._settings.premium_stars_price + return ( + "Стоимость генерации:\n" + f"• обычная очередь — {normal} звёзд;\n" + f"• премиум-очередь — {premium} звёзд.\n\n" + + START_TEXT + ) return START_TEXT async def _start( @@ -156,11 +172,13 @@ class MpakoBot: await message.reply_text("Активных задач нет.") return if active.status == "running": - await message.reply_text("Твоя картинка сейчас генерируется.") + label = " (премиум)" if active.is_premium else "" + await message.reply_text(f"Твоя картинка сейчас генерируется{label}.") return position = await self._queue.position(active.id) if self._queue else 0 + label = "премиум-очередь" if active.is_premium else "очередь" await message.reply_text( - f"Задача в очереди. Перед тобой: {position} {self._people(position)}." + f"Задача в {label}. Перед тобой: {position} {self._people(position)}." ) async def _pre_checkout( @@ -187,6 +205,74 @@ class MpakoBot: except TelegramError: logger.exception("Не удалось ответить на pre_checkout_query") + async def _queue_choice( + self, update: Update, context: ContextTypes.DEFAULT_TYPE + ) -> None: + query = update.callback_query + if query is None or query.data is None: + return + try: + _, payload, tier = query.data.split(":", 2) + except ValueError: + await query.answer("Некорректный выбор.", show_alert=True) + return + if tier not in {"normal", "premium"}: + await query.answer("Некорректный выбор.", show_alert=True) + return + + pending = await self._db.get_pending_payment(payload) + if ( + pending is None + or pending.user_id != query.from_user.id + or pending.status != "pending" + or pending.amount != 0 + ): + await query.answer("Счёт уже создан или устарел.", show_alert=True) + return + + active = await self._db.get_active_by_user(query.from_user.id) + if active is not None: + await query.answer( + "У тебя уже есть активная задача.", show_alert=True + ) + return + + is_premium = tier == "premium" + amount = ( + self._settings.premium_stars_price + if is_premium + else self._settings.stars_price + ) + if not await self._db.set_pending_choice(payload, is_premium, amount): + await query.answer("Счёт уже создан.", show_alert=True) + return + await query.answer() + + label = "Премиум-очередь" if is_premium else "Обычная очередь" + try: + await self._retry( + self._bot.send_invoice, + chat_id=query.from_user.id, + title=label, + description=( + "Приоритетная генерация." if is_premium + else "Обычная генерация." + ), + payload=payload, + provider_token="", + currency="XTR", + prices=[LabeledPrice(label=label, amount=amount)], + ) + except TelegramError: + logger.exception("Не удалось выставить счёт пользователю %s", query.from_user.id) + await self._db.mark_pending_cancelled(payload) + await query.edit_message_text("Не удалось выставить счёт. Попробуй позже.") + return + + await query.edit_message_text( + f"Счёт на {amount} звёзд отправлен. Очередь: {label.lower()}." + ) + async def _successful_payment( self, update: Update, context: ContextTypes.DEFAULT_TYPE ) -> None: @@ -231,6 +317,7 @@ class MpakoBot: skip_payment=True, payment_charge_id=charge_id, payment_payload=payload, + is_premium=pending.is_premium, ) async def _text( @@ -376,6 +463,7 @@ class MpakoBot: skip_payment: bool = False, payment_charge_id: str | None = None, payment_payload: str | None = None, + is_premium: bool = False, ) -> bool: message = update.effective_message user = update.effective_user @@ -414,8 +502,18 @@ class MpakoBot: await message.reply_text(BLOCKED_TEXT) return False + active = await self._db.get_active_by_user(user.id) + if active is not None: + if active.status == "running": + text = "Твоя картинка уже генерируется." + else: + ahead = await self._queue.position(active.id) + text = f"Ты уже в очереди. Перед тобой: {ahead} {self._people(ahead)}." + await message.reply_text(text) + return False + + counts = await self._queue.queue_counts() payload = uuid.uuid4().hex - amount = self._settings.stars_price await self._db.create_pending_payment( payload=payload, user_id=user.id, @@ -423,28 +521,44 @@ class MpakoBot: prompt=prompt, username=user.username, input_file_path=input_path, - amount=amount, + amount=0, + ) + keyboard = InlineKeyboardMarkup( + [ + [ + InlineKeyboardButton( + f"Обычная — {self._settings.stars_price} звёзд", + callback_data=f"queue:{payload}:normal", + ) + ], + [ + InlineKeyboardButton( + f"Премиум — {self._settings.premium_stars_price} звёзд", + callback_data=f"queue:{payload}:premium", + ) + ], + ] + ) + text = ( + f"Обычная очередь: {counts['normal']} чел.\n" + f"Премиум-очередь: {counts['premium']} чел.\n\n" + "Выбери очередь:" ) try: await self._retry( - self._bot.send_invoice, - chat_id=user.id, - title="Генерация изображения", - description="Одна генерация или редактирование изображения.", - payload=payload, - provider_token="", - currency="XTR", - prices=[LabeledPrice(label="Генерация", amount=amount)], + message.reply_text, + text, + reply_markup=keyboard, ) except TelegramError: - logger.exception("Не удалось выставить счёт пользователю %s", user.id) + logger.exception("Не удалось показать выбор очереди %s", user.id) await self._db.mark_pending_cancelled(payload) - await message.reply_text( - "Не удалось выставить счёт. Попробуй позже." - ) return False return True + if is_admin: + is_premium = True + try: request, position = await self._queue.enqueue( user_id=user.id, @@ -455,6 +569,7 @@ class MpakoBot: has_person=has_person, is_admin=is_admin, payment_charge_id=payment_charge_id, + is_premium=is_premium, ) except PromptBlocked: note = await self._refund_if_paid( @@ -497,7 +612,7 @@ class MpakoBot: try: status_message = await self._retry( - message.reply_text, self._ack_text(position) + message.reply_text, self._ack_text(position, is_premium) ) except TelegramError: logger.exception( @@ -553,10 +668,11 @@ class MpakoBot: logger.exception("Не удалось вернуть звёзды пользователю %s", user_id) return False - def _ack_text(self, position: int) -> str: + def _ack_text(self, position: int, is_premium: bool = False) -> str: + label = "Премиум-очередь. " if is_premium else "" if position <= 0: - return "Принято. Ты первый, начинаю…" - return f"Принято. Перед тобой: {position} {self._people(position)}." + return f"Принято. {label}Ты первый, начинаю…" + return f"Принято. {label}Перед тобой: {position} {self._people(position)}." @staticmethod def _people(count: int) -> str: diff --git a/checklist.md b/checklist.md index aa89889..1d8caaa 100644 --- a/checklist.md +++ b/checklist.md @@ -8,7 +8,7 @@ - Реализованы блоки 2–7, базовое хранилище (9), файлы Docker (10), тесты (11), README (12), веб-история (13), улучшение промптов (14), модерация (15) и оплата Stars (16). - Вживую проверены: T2I (Z-Image), I2I (Krea2), WebSocket-прогресс, улучшение промптов, Docker, CI и веб-история. -- 46 тестов проходят: БД, workflows, очередь, веб-история, enhancer, модерация, платежи, вспомогательные функции ComfyUI. +- 49 тестов проходят: БД, workflows, очередь, веб-история, enhancer, модерация, платежи, премиум-очередь, вспомогательные функции ComfyUI. - Не проверено: реальная оплата Stars в Telegram, очередь из 3+ пользователей, падение ComfyUI, рестарт с `running`-задачей, детекция лица на реальном фото. ## 0. Допущения и решения по умолчанию @@ -332,6 +332,7 @@ ON requests(user_id) WHERE status IN ('queued', 'running'); ## 16. Оплата Stars - [x] `PAYMENTS_ENABLED` / `STARS_PRICE`, админы бесплатно. +- [x] Премиум-очередь: `PREMIUM_STARS_PRICE`, чередование с обычной, счётчики очередей до оплаты. - [x] Invoice `XTR` + `pre_checkout_query` + `successful_payment`. - [x] `pending_payments`: payload, сумма, статус, charge_id, request_id. - [x] Возврат звёзд при ошибке постановки в очередь и генерации. diff --git a/config.py b/config.py index 10bf059..46d21cd 100644 --- a/config.py +++ b/config.py @@ -96,6 +96,7 @@ class Settings: blocked_public_figures: frozenset[str] = field(default_factory=frozenset) payments_enabled: bool = True stars_price: int = 5 + premium_stars_price: int = 15 web_page_size: int = 50 web_bind_ip: str = "127.0.0.1" @@ -156,6 +157,7 @@ class Settings: blocked_public_figures=_csv("BLOCKED_PUBLIC_FIGURES"), payments_enabled=_bool("PAYMENTS_ENABLED", True), stars_price=max(1, _int("STARS_PRICE", 5)), + premium_stars_price=max(1, _int("PREMIUM_STARS_PRICE", 15)), web_page_size=max(10, min(200, _int("WEB_PAGE_SIZE", 50))), web_bind_ip=os.getenv("WEB_BIND_IP", "127.0.0.1").strip() or "127.0.0.1", diff --git a/db.py b/db.py index 7cbbf82..e8034b8 100644 --- a/db.py +++ b/db.py @@ -39,6 +39,7 @@ class Request: comfy_prompt_id: str | None payment_charge_id: str | None payment_refunded: bool + is_premium: bool mode: str seed: int | None error_text: str | None @@ -63,6 +64,7 @@ class Request: comfy_prompt_id=row["comfy_prompt_id"], payment_charge_id=row["payment_charge_id"], payment_refunded=bool(row["payment_refunded"]), + is_premium=bool(row["is_premium"]), mode=row["mode"], seed=row["seed"], error_text=row["error_text"], @@ -116,6 +118,7 @@ class PendingPayment: input_file_path: str | None amount: int status: str + is_premium: bool telegram_payment_charge_id: str | None request_id: int | None @@ -131,6 +134,7 @@ class PendingPayment: input_file_path=row["input_file_path"], amount=int(row["amount"]), status=row["status"], + is_premium=bool(row["is_premium"]), telegram_payment_charge_id=row["telegram_payment_charge_id"], request_id=row["request_id"], ) @@ -151,6 +155,7 @@ CREATE TABLE IF NOT EXISTS requests ( comfy_prompt_id TEXT, payment_charge_id TEXT, payment_refunded INTEGER NOT NULL DEFAULT 0, + is_premium INTEGER NOT NULL DEFAULT 0, mode TEXT NOT NULL CHECK (mode IN ('t2i', 'i2i')), seed INTEGER, error_text TEXT, @@ -190,6 +195,7 @@ CREATE TABLE IF NOT EXISTS pending_payments ( amount INTEGER NOT NULL, status TEXT NOT NULL DEFAULT 'pending' CHECK (status IN ('pending', 'paid', 'cancelled', 'refunded')), + is_premium INTEGER NOT NULL DEFAULT 0, telegram_payment_charge_id TEXT, request_id INTEGER ); @@ -234,6 +240,16 @@ class Database: await conn.execute( "ALTER TABLE requests ADD COLUMN payment_refunded INTEGER NOT NULL DEFAULT 0" ) + if "is_premium" not in columns: + await conn.execute( + "ALTER TABLE requests ADD COLUMN is_premium INTEGER NOT NULL DEFAULT 0" + ) + async with conn.execute("PRAGMA table_info(pending_payments)") as cursor: + pending_columns = {row["name"] async for row in cursor} + if "is_premium" not in pending_columns: + await conn.execute( + "ALTER TABLE pending_payments ADD COLUMN is_premium INTEGER NOT NULL DEFAULT 0" + ) async def close(self) -> None: if self._conn is not None: @@ -255,6 +271,7 @@ class Database: input_file_path: str | None, seed: int, payment_charge_id: str | None = None, + is_premium: bool = False, ) -> Request: conn = self._connection() async with self._write_lock: @@ -263,8 +280,9 @@ class Database: """ INSERT INTO requests (user_id, username, timestamp, prompt_text, - input_file_path, status, mode, seed, payment_charge_id) - VALUES (?, ?, ?, ?, ?, 'queued', ?, ?, ?) + input_file_path, status, mode, seed, payment_charge_id, + is_premium) + VALUES (?, ?, ?, ?, ?, 'queued', ?, ?, ?, ?) """, ( user_id, @@ -275,6 +293,7 @@ class Database: mode, seed, payment_charge_id, + int(is_premium), ), ) request_id = cursor.lastrowid @@ -418,16 +437,42 @@ class Database: async def count_ahead(self, request_id: int) -> int: conn = self._connection() - async with conn.execute( + request = await self.get_request(request_id) + if request is None: + return 0 + if request.is_premium: + query = """ + SELECT COUNT(*) AS total FROM requests + WHERE is_premium = 1 AND status IN ('queued', 'running') + AND id < ? """ - SELECT COUNT(*) AS total FROM requests - WHERE id < ? AND status IN ('queued', 'running') - """, - (request_id,), - ) as cursor: + params = (request_id,) + else: + query = """ + SELECT COUNT(*) AS total FROM requests + WHERE status IN ('queued', 'running') + AND (is_premium = 1 OR (is_premium = 0 AND id < ?)) + """ + params = (request_id,) + async with conn.execute(query, params) as cursor: row = await cursor.fetchone() return int(row["total"]) if row is not None else 0 + async def queue_counts(self) -> dict[str, int]: + conn = self._connection() + counts = {"premium": 0, "normal": 0} + async with conn.execute( + """ + SELECT is_premium, COUNT(*) AS total FROM requests + WHERE status IN ('queued', 'running') + GROUP BY is_premium + """ + ) as cursor: + async for row in cursor: + key = "premium" if row["is_premium"] else "normal" + counts[key] = int(row["total"]) + return counts + async def set_running(self, request_id: int, comfy_prompt_id: str | None = None) -> None: await self._update( """ @@ -468,6 +513,7 @@ class Database: username: str | None, input_file_path: str | None, amount: int, + is_premium: bool = False, ) -> PendingPayment: conn = self._connection() async with self._write_lock: @@ -475,8 +521,8 @@ class Database: """ INSERT INTO pending_payments (payload, user_id, timestamp, mode, prompt, username, - input_file_path, amount, status) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, 'pending') + input_file_path, amount, status, is_premium) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, 'pending', ?) """, ( payload, @@ -487,6 +533,7 @@ class Database: username, input_file_path, amount, + int(is_premium), ), ) await conn.commit() @@ -503,6 +550,22 @@ class Database: row = await cursor.fetchone() return PendingPayment.from_row(row) if row is not None else None + async def set_pending_choice( + self, payload: str, is_premium: bool, amount: int + ) -> bool: + conn = self._connection() + async with self._write_lock: + cursor = await conn.execute( + """ + UPDATE pending_payments + SET is_premium = ?, amount = ? + WHERE payload = ? AND status = 'pending' AND amount = 0 + """, + (int(is_premium), amount, payload), + ) + await conn.commit() + return cursor.rowcount > 0 + async def mark_pending_paid(self, payload: str, charge_id: str) -> bool: conn = self._connection() async with self._write_lock: diff --git a/queue_manager.py b/queue_manager.py index d7a60c0..79dbe52 100644 --- a/queue_manager.py +++ b/queue_manager.py @@ -61,7 +61,9 @@ class QueueManager: self._on_event = on_event self._enhancer = enhancer - self._queue: deque[int] = deque() + self._premium: deque[int] = deque() + self._normal: deque[int] = deque() + self._next_premium = True self._active_users: set[int] = set() self._wake = asyncio.Event() self._lock = asyncio.Lock() @@ -74,7 +76,8 @@ class QueueManager: recovered = await self._db.recover() async with self._lock: for request in recovered: - self._queue.append(request.id) + target = self._premium if request.is_premium else self._normal + target.append(request.id) self._active_users.add(request.user_id) if recovered: logger.info("Восстановлено из очереди: %s", len(recovered)) @@ -108,6 +111,7 @@ class QueueManager: has_person: bool = False, is_admin: bool = False, payment_charge_id: str | None = None, + is_premium: bool = False, ) -> tuple[Request, int]: async with self._lock: if self._stopping: @@ -131,7 +135,7 @@ class QueueManager: active = await self._db.get_active_by_user(user_id) if active is not None: raise UserAlreadyQueued(user_id) - if len(self._queue) >= self._settings.max_queue_size: + if len(self._premium) + len(self._normal) >= self._settings.max_queue_size: raise QueueFull() seed = secrets.randbelow(2**32) @@ -143,15 +147,56 @@ class QueueManager: input_file_path=input_file_path, seed=seed, payment_charge_id=payment_charge_id, + is_premium=is_premium, ) - self._queue.append(request.id) + target = self._premium if is_premium else self._normal + target.append(request.id) self._active_users.add(user_id) self._wake.set() position = await self._db.count_ahead(request.id) return request, position + async def queue_counts(self) -> dict[str, int]: + return await self._db.queue_counts() + async def position(self, request_id: int) -> int: - return await self._db.count_ahead(request_id) + request = await self._db.get_request(request_id) + if request is not None and request.status == "running": + return 0 + + counts = await self._db.queue_counts() + running = max( + 0, + sum(counts.values()) - len(self._premium) - len(self._normal), + ) + premium = list(self._premium) + normal = list(self._normal) + prefer_premium = self._next_premium + premium_index = 0 + normal_index = 0 + ahead = running + while premium_index < len(premium) or normal_index < len(normal): + if prefer_premium and premium_index < len(premium): + current = premium[premium_index] + premium_index += 1 + taken_premium = True + elif not prefer_premium and normal_index < len(normal): + current = normal[normal_index] + normal_index += 1 + taken_premium = False + elif premium_index < len(premium): + current = premium[premium_index] + premium_index += 1 + taken_premium = True + else: + current = normal[normal_index] + normal_index += 1 + taken_premium = False + if current == request_id: + return ahead + ahead += 1 + prefer_premium = not taken_premium + return ahead async def _worker(self, index: int) -> None: logger.info("Воркер %s запущен", index) @@ -166,12 +211,25 @@ class QueueManager: async def _next_request(self) -> int | None: while not self._stopping: - if self._queue: - return self._queue.popleft() + if self._next_premium: + if self._premium: + return self._take(self._premium, premium=True) + if self._normal: + return self._take(self._normal, premium=False) + else: + if self._normal: + return self._take(self._normal, premium=False) + if self._premium: + return self._take(self._premium, premium=True) self._wake.clear() await self._wake.wait() return None + def _take(self, queue: deque[int], *, premium: bool) -> int: + request_id = queue.popleft() + self._next_premium = not premium + return request_id + async def _process(self, request_id: int) -> None: request = await self._db.get_request(request_id) if request is None: diff --git a/tests/test_db.py b/tests/test_db.py index 23e5a87..bea616a 100644 --- a/tests/test_db.py +++ b/tests/test_db.py @@ -165,6 +165,32 @@ async def test_history_filters_and_stats(db: Database) -> None: 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", @@ -173,10 +199,16 @@ async def test_pending_payment_lifecycle(db: Database) -> None: prompt="apple", username="alice", input_file_path=None, - amount=5, + amount=0, ) assert pending.status == "pending" - assert pending.amount == 5 + 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") diff --git a/tests/test_queue.py b/tests/test_queue.py index a036132..1a53984 100644 --- a/tests/test_queue.py +++ b/tests/test_queue.py @@ -251,3 +251,46 @@ async def test_admin_prompt_is_not_filtered(db: Database, tmp_path: Path) -> Non assert await db.list_blocked() == [] await queue.stop() + + +async def test_premium_queue_goes_first(db: Database, tmp_path: Path) -> None: + events: list[QueueEvent] = [] + queue = make_queue(db, make_settings(tmp_path), events) + + normal, _ = await queue.enqueue( + user_id=1, mode="t2i", prompt="normal" + ) + premium, _ = await queue.enqueue( + user_id=2, mode="t2i", prompt="premium", is_premium=True + ) + + first = await queue._next_request() # type: ignore[attr-defined] + second = await queue._next_request() # type: ignore[attr-defined] + assert first == premium.id + assert second == normal.id + + await queue.stop() + + +async def test_premium_and_normal_alternate(db: Database, tmp_path: Path) -> None: + events: list[QueueEvent] = [] + queue = make_queue(db, make_settings(tmp_path), events) + + normal_1, _ = await queue.enqueue(user_id=1, mode="t2i", prompt="n1") + normal_2, _ = await queue.enqueue(user_id=2, mode="t2i", prompt="n2") + premium_1, _ = await queue.enqueue( + user_id=3, mode="t2i", prompt="p1", is_premium=True + ) + premium_2, _ = await queue.enqueue( + user_id=4, mode="t2i", prompt="p2", is_premium=True + ) + + order = [await queue._next_request() for _ in range(4)] # type: ignore[attr-defined] + assert order == [ + premium_1.id, + normal_1.id, + premium_2.id, + normal_2.id, + ] + + await queue.stop() diff --git a/web/templates/detail.html b/web/templates/detail.html index 0ef49eb..bf9bb0b 100644 --- a/web/templates/detail.html +++ b/web/templates/detail.html @@ -11,6 +11,8 @@
{{ item.status }}
Режим
{{ item.mode }}
+
Очередь
+
{% if item.is_premium %}премиум{% else %}обычная{% endif %}
Создан
{{ item.timestamp|fmt_dt }}
Завершён
diff --git a/web/templates/index.html b/web/templates/index.html index 20632b9..b5f1518 100644 --- a/web/templates/index.html +++ b/web/templates/index.html @@ -77,7 +77,7 @@ {{ item.timestamp|fmt_dt }} {{ item.display_user }} - {{ item.mode }} + {{ item.mode }}{% if item.is_premium %} P{% endif %} {{ item.status }} {% if item.enhanced_prompt %}