From ee173190a05764ab0a735dbbc30bb50257f34e9d Mon Sep 17 00:00:00 2001 From: Egor Date: Sun, 9 Nov 2025 05:48:45 +0300 Subject: [PATCH] Fix category edit menu callback --- app/database/crud/server_category.py | 57 ++++ app/database/crud/server_squad.py | 87 +++++- app/database/models.py | 24 +- app/database/universal_migration.py | 122 +++++++++ app/handlers/admin/servers.py | 250 +++++++++++++++++- app/services/backup_service.py | 5 +- app/services/subscription_purchase_service.py | 172 +++++++++++- app/states.py | 1 + ...test_subscription_auto_purchase_service.py | 111 ++++++++ 9 files changed, 805 insertions(+), 24 deletions(-) create mode 100644 app/database/crud/server_category.py diff --git a/app/database/crud/server_category.py b/app/database/crud/server_category.py new file mode 100644 index 00000000..8e81c9c8 --- /dev/null +++ b/app/database/crud/server_category.py @@ -0,0 +1,57 @@ +from typing import List, Optional + +from sqlalchemy import select, update, delete +from sqlalchemy.ext.asyncio import AsyncSession + +from app.database.models import ServerCategory + + +async def get_all_server_categories(db: AsyncSession, include_inactive: bool = False) -> List[ServerCategory]: + query = select(ServerCategory).order_by(ServerCategory.sort_order, ServerCategory.name) + if not include_inactive: + query = query.where(ServerCategory.is_active.is_(True)) + result = await db.execute(query) + return result.scalars().all() + + +async def get_server_category_by_id(db: AsyncSession, category_id: int) -> Optional[ServerCategory]: + result = await db.execute( + select(ServerCategory).where(ServerCategory.id == category_id) + ) + return result.scalar_one_or_none() + + +async def create_server_category( + db: AsyncSession, + name: str, + sort_order: int = 0, + is_active: bool = True, +) -> ServerCategory: + category = ServerCategory(name=name, sort_order=sort_order, is_active=is_active) + db.add(category) + await db.commit() + await db.refresh(category) + return category + + +async def update_server_category( + db: AsyncSession, + category_id: int, + **updates, +) -> Optional[ServerCategory]: + valid_fields = {"name", "sort_order", "is_active"} + filtered_updates = {k: v for k, v in updates.items() if k in valid_fields} + if not filtered_updates: + return await get_server_category_by_id(db, category_id) + + await db.execute( + update(ServerCategory).where(ServerCategory.id == category_id).values(**filtered_updates) + ) + await db.commit() + return await get_server_category_by_id(db, category_id) + + +async def delete_server_category(db: AsyncSession, category_id: int) -> bool: + await db.execute(delete(ServerCategory).where(ServerCategory.id == category_id)) + await db.commit() + return True diff --git a/app/database/crud/server_squad.py b/app/database/crud/server_squad.py index c08a8676..5af00af6 100644 --- a/app/database/crud/server_squad.py +++ b/app/database/crud/server_squad.py @@ -19,6 +19,7 @@ from sqlalchemy.orm import selectinload from app.database.models import ( PromoGroup, + ServerCategory, ServerSquad, SubscriptionServer, Subscription, @@ -48,6 +49,7 @@ async def create_server_squad( is_available: bool = True, is_trial_eligible: bool = False, sort_order: int = 0, + category_id: Optional[int] = None, promo_group_ids: Optional[Iterable[int]] = None, ) -> ServerSquad: @@ -71,6 +73,13 @@ async def create_server_squad( "Не все промогруппы найдены при создании сервера %s", display_name ) + if category_id is not None: + category_exists = await db.execute( + select(ServerCategory.id).where(ServerCategory.id == category_id) + ) + if category_exists.scalar_one_or_none() is None: + raise ValueError("Server category not found") + server_squad = ServerSquad( squad_uuid=squad_uuid, display_name=display_name, @@ -82,6 +91,7 @@ async def create_server_squad( is_available=is_available, is_trial_eligible=is_trial_eligible, sort_order=sort_order, + category_id=category_id, allowed_promo_groups=promo_groups, ) @@ -100,7 +110,10 @@ async def get_server_squad_by_uuid( result = await db.execute( select(ServerSquad) - .options(selectinload(ServerSquad.allowed_promo_groups)) + .options( + selectinload(ServerSquad.allowed_promo_groups), + selectinload(ServerSquad.category), + ) .where(ServerSquad.squad_uuid == squad_uuid) ) return result.scalars().unique().one_or_none() @@ -113,7 +126,10 @@ async def get_server_squad_by_id( result = await db.execute( select(ServerSquad) - .options(selectinload(ServerSquad.allowed_promo_groups)) + .options( + selectinload(ServerSquad.allowed_promo_groups), + selectinload(ServerSquad.category), + ) .where(ServerSquad.id == server_id) ) return result.scalars().unique().one_or_none() @@ -126,8 +142,8 @@ async def get_all_server_squads( limit: int = 50 ) -> Tuple[List[ServerSquad], int]: - query = select(ServerSquad) - + query = select(ServerSquad).options(selectinload(ServerSquad.category)) + if available_only: query = query.where(ServerSquad.is_available == True) @@ -151,11 +167,16 @@ async def get_all_server_squads( async def get_available_server_squads( db: AsyncSession, promo_group_id: Optional[int] = None, + category_id: Optional[int] = None, + only_with_capacity: bool = False, ) -> List[ServerSquad]: query = ( select(ServerSquad) - .options(selectinload(ServerSquad.allowed_promo_groups)) + .options( + selectinload(ServerSquad.allowed_promo_groups), + selectinload(ServerSquad.category), + ) .where(ServerSquad.is_available.is_(True)) .order_by(ServerSquad.sort_order, ServerSquad.display_name) ) @@ -165,8 +186,21 @@ async def get_available_server_squads( PromoGroup.id == promo_group_id ) + if category_id is not None: + query = query.where(ServerSquad.category_id == category_id) + result = await db.execute(query) - return result.scalars().unique().all() + squads = result.scalars().unique().all() + + if only_with_capacity: + filtered: List[ServerSquad] = [] + for squad in squads: + if squad.is_full: + continue + filtered.append(squad) + return filtered + + return squads async def get_active_server_squads(db: AsyncSession) -> List[ServerSquad]: @@ -221,6 +255,35 @@ async def get_random_active_squad_uuid( return fallback_uuid +async def choose_least_loaded_server_in_category( + db: AsyncSession, + category_id: int, + promo_group_id: Optional[int] = None, +) -> Optional[ServerSquad]: + """Выбирает наименее загруженный сервер в категории.""" + + category_squads = await get_available_server_squads( + db, + promo_group_id=promo_group_id, + category_id=category_id, + only_with_capacity=True, + ) + + if not category_squads: + return None + + def load_key(squad: ServerSquad) -> tuple: + max_users = squad.max_users or 0 + current_users = squad.current_users or 0 + if max_users: + ratio = current_users / max_users + else: + ratio = 0 + return (ratio, current_users, squad.id) + + return min(category_squads, key=load_key) + + async def update_server_squad_promo_groups( db: AsyncSession, server_id: int, promo_group_ids: Iterable[int] ) -> Optional[ServerSquad]: @@ -271,10 +334,20 @@ async def update_server_squad( "is_available", "sort_order", "is_trial_eligible", + "category_id", } filtered_updates = {k: v for k, v in updates.items() if k in valid_fields} - + + if "category_id" in filtered_updates: + category_id = filtered_updates["category_id"] + if category_id is not None: + category_exists = await db.execute( + select(ServerCategory.id).where(ServerCategory.id == category_id) + ) + if category_exists.scalar_one_or_none() is None: + raise ValueError("Server category not found") + if not filtered_updates: return None diff --git a/app/database/models.py b/app/database/models.py index 09695965..e0fb5cb8 100644 --- a/app/database/models.py +++ b/app/database/models.py @@ -1281,6 +1281,22 @@ class PollAnswer(Base): ) +class ServerCategory(Base): + __tablename__ = "server_categories" + + id = Column(Integer, primary_key=True, index=True) + name = Column(String(255), nullable=False, unique=True) + sort_order = Column(Integer, default=0) + is_active = Column(Boolean, default=True, nullable=False) + created_at = Column(DateTime, default=func.now(), nullable=False) + updated_at = Column(DateTime, default=func.now(), onupdate=func.now(), nullable=False) + + servers = relationship("ServerSquad", back_populates="category") + + def __repr__(self) -> str: # pragma: no cover - debug helper + return f"" + + class ServerSquad(Base): __tablename__ = "server_squads" @@ -1302,8 +1318,10 @@ class ServerSquad(Base): description = Column(Text, nullable=True) sort_order = Column(Integer, default=0) - - max_users = Column(Integer, nullable=True) + + category_id = Column(Integer, ForeignKey("server_categories.id", ondelete="SET NULL"), nullable=True) + + max_users = Column(Integer, nullable=True) current_users = Column(Integer, default=0) created_at = Column(DateTime, default=func.now()) @@ -1315,6 +1333,8 @@ class ServerSquad(Base): back_populates="server_squads", lazy="selectin", ) + + category = relationship("ServerCategory", back_populates="servers") @property def price_rubles(self) -> float: diff --git a/app/database/universal_migration.py b/app/database/universal_migration.py index 50f05d5b..88e062bc 100644 --- a/app/database/universal_migration.py +++ b/app/database/universal_migration.py @@ -3053,6 +3053,63 @@ async def ensure_server_promo_groups_setup() -> bool: return False +async def ensure_server_categories_table() -> bool: + logger.info("=== ПРОВЕРКА ТАБЛИЦЫ КАТЕГОРИЙ СЕРВЕРОВ ===") + + try: + table_exists = await check_table_exists("server_categories") + + if table_exists: + logger.info("ℹ️ Таблица server_categories уже существует") + return True + + async with engine.begin() as conn: + db_type = await get_database_type() + + if db_type == "sqlite": + create_sql = """ + CREATE TABLE server_categories ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + name VARCHAR(255) NOT NULL UNIQUE, + sort_order INTEGER NOT NULL DEFAULT 0, + is_active BOOLEAN NOT NULL DEFAULT 1, + created_at DATETIME DEFAULT CURRENT_TIMESTAMP, + updated_at DATETIME DEFAULT CURRENT_TIMESTAMP + ); + """ + elif db_type == "postgresql": + create_sql = """ + CREATE TABLE server_categories ( + id SERIAL PRIMARY KEY, + name VARCHAR(255) NOT NULL UNIQUE, + sort_order INTEGER NOT NULL DEFAULT 0, + is_active BOOLEAN NOT NULL DEFAULT TRUE, + created_at TIMESTAMP DEFAULT NOW(), + updated_at TIMESTAMP DEFAULT NOW() + ); + """ + else: + create_sql = """ + CREATE TABLE server_categories ( + id INT AUTO_INCREMENT PRIMARY KEY, + name VARCHAR(255) NOT NULL UNIQUE, + sort_order INT NOT NULL DEFAULT 0, + is_active BOOLEAN NOT NULL DEFAULT TRUE, + created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, + updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP + ); + """ + + await conn.execute(text(create_sql)) + logger.info("✅ Таблица server_categories создана") + + return True + + except Exception as error: + logger.error(f"Ошибка создания таблицы server_categories: {error}") + return False + + async def add_server_trial_flag_column() -> bool: column_exists = await check_column_exists('server_squads', 'is_trial_eligible') if column_exists: @@ -3087,6 +3144,53 @@ async def add_server_trial_flag_column() -> bool: return False +async def add_server_category_column() -> bool: + column_exists = await check_column_exists('server_squads', 'category_id') + if column_exists: + logger.info("Колонка category_id уже существует в server_squads") + return True + + try: + async with engine.begin() as conn: + db_type = await get_database_type() + + column_def = 'INTEGER' + + await conn.execute( + text(f"ALTER TABLE server_squads ADD COLUMN category_id {column_def}") + ) + + if db_type == 'postgresql': + await conn.execute( + text( + "ALTER TABLE server_squads ADD CONSTRAINT " + "fk_server_squads_category FOREIGN KEY (category_id) " + "REFERENCES server_categories(id) ON DELETE SET NULL" + ) + ) + elif db_type != 'sqlite': + await conn.execute( + text( + "ALTER TABLE server_squads ADD CONSTRAINT fk_server_squads_category " + "FOREIGN KEY (category_id) REFERENCES server_categories(id) ON DELETE SET NULL" + ) + ) + + await conn.execute( + text( + "CREATE INDEX IF NOT EXISTS idx_server_squads_category_id " + "ON server_squads(category_id)" + ) + ) + + logger.info("✅ Добавлена колонка category_id в server_squads") + return True + + except Exception as error: + logger.error(f"Ошибка добавления колонки category_id: {error}") + return False + + async def create_system_settings_table() -> bool: table_exists = await check_table_exists("system_settings") if table_exists: @@ -4039,6 +4143,18 @@ async def run_universal_migration(): else: logger.warning("⚠️ Проблемы с настройкой доступа серверов к промогруппам") + categories_ready = await ensure_server_categories_table() + if categories_ready: + logger.info("✅ Таблица категорий серверов готова") + else: + logger.warning("⚠️ Проблемы с созданием таблицы категорий серверов") + + category_column_ready = await add_server_category_column() + if category_column_ready: + logger.info("✅ Колонка категории для серверов настроена") + else: + logger.warning("⚠️ Проблемы с добавлением колонки категории для серверов") + logger.info("=== ОБНОВЛЕНИЕ ВНЕШНИХ КЛЮЧЕЙ ===") fk_updated = await fix_foreign_keys_for_user_deletion() if fk_updated: @@ -4115,6 +4231,8 @@ async def check_migration_status(): "promo_groups_table": False, "server_promo_groups_table": False, "server_squads_trial_column": False, + "server_categories_table": False, + "server_squads_category_column": False, "privacy_policies_table": False, "public_offers_table": False, "users_promo_group_column": False, @@ -4148,6 +4266,8 @@ async def check_migration_status(): status["promo_groups_table"] = await check_table_exists('promo_groups') status["server_promo_groups_table"] = await check_table_exists('server_squad_promo_groups') status["server_squads_trial_column"] = await check_column_exists('server_squads', 'is_trial_eligible') + status["server_categories_table"] = await check_table_exists('server_categories') + status["server_squads_category_column"] = await check_column_exists('server_squads', 'category_id') status["discount_offers_table"] = await check_table_exists('discount_offers') status["discount_offers_effect_column"] = await check_column_exists('discount_offers', 'effect_type') @@ -4204,6 +4324,8 @@ async def check_migration_status(): "promo_groups_table": "Таблица промо-групп", "server_promo_groups_table": "Связи серверов и промогрупп", "server_squads_trial_column": "Колонка триального назначения у серверов", + "server_categories_table": "Таблица категорий серверов", + "server_squads_category_column": "Колонка категории у серверов", "users_promo_group_column": "Колонка promo_group_id у пользователей", "promo_groups_period_discounts_column": "Колонка period_discounts у промо-групп", "promo_groups_auto_assign_column": "Колонка auto_assign_total_spent_kopeks у промо-групп", diff --git a/app/handlers/admin/servers.py b/app/handlers/admin/servers.py index 7ccb20cc..6b5b01de 100644 --- a/app/handlers/admin/servers.py +++ b/app/handlers/admin/servers.py @@ -1,5 +1,7 @@ import html import logging +from typing import List + from aiogram import Dispatcher, types, F from aiogram.fsm.context import FSMContext from sqlalchemy.ext.asyncio import AsyncSession @@ -18,6 +20,10 @@ from app.database.crud.server_squad import ( update_server_squad_promo_groups, get_server_connected_users, ) +from app.database.crud.server_category import ( + get_all_server_categories, + create_server_category, +) from app.database.crud.promo_group import get_promo_groups_with_counts from app.services.remnawave_service import RemnaWaveService from app.utils.decorators import admin_required, error_handler @@ -36,6 +42,7 @@ def _build_server_edit_view(server): ) trial_status = "✅ Да" if server.is_trial_eligible else "⚪️ Нет" + category_name = getattr(getattr(server, "category", None), "name", None) or "Не указана" text = f""" 🌐 Редактирование сервера @@ -53,6 +60,7 @@ def _build_server_edit_view(server): • Лимит пользователей: {server.max_users or 'Без лимита'} • Текущих пользователей: {server.current_users} • Промогруппы: {promo_groups_text} +• Категория: {category_name} • Выдача триала: {trial_status} Описание: @@ -97,6 +105,11 @@ def _build_server_edit_view(server): text="📝 Описание", callback_data=f"admin_server_edit_desc_{server.id}" ), ], + [ + types.InlineKeyboardButton( + text="🏷 Категория", callback_data=f"admin_server_edit_category_{server.id}" + ), + ], [ types.InlineKeyboardButton( text="❌ Отключить" if server.is_available else "✅ Включить", @@ -1074,6 +1087,225 @@ async def process_server_description_edit( await message.answer("❌ Ошибка при обновлении сервера") +@admin_required +@error_handler +async def show_server_category_menu( + callback: types.CallbackQuery, + db_user: User, + db: AsyncSession, +): + server_id = int(callback.data.split("_")[-1]) + server = await get_server_squad_by_id(db, server_id) + + if not server: + await callback.answer("❌ Сервер не найден!", show_alert=True) + return + + categories = await get_all_server_categories(db) + + current_category = getattr(getattr(server, "category", None), "name", None) + safe_current = html.escape(current_category or "Не выбрана") + + text_lines = [ + "🏷 Категория сервера", + "", + f"Текущая категория: {safe_current}", + "", + "Выберите категорию из списка или создайте новую:", + ] + + keyboard_rows: List[List[types.InlineKeyboardButton]] = [] + + for category in categories: + emoji = "✅" if server.category_id == category.id else "⚪️" + keyboard_rows.append( + [ + types.InlineKeyboardButton( + text=f"{emoji} {category.name}", + callback_data=f"admin_server_category_set_{server.id}_{category.id}", + ) + ] + ) + + keyboard_rows.append( + [ + types.InlineKeyboardButton( + text="➕ Создать категорию", + callback_data=f"admin_server_category_create_{server.id}", + ) + ] + ) + + if server.category_id: + keyboard_rows.append( + [ + types.InlineKeyboardButton( + text="🗑 Удалить категорию", + callback_data=f"admin_server_category_clear_{server.id}", + ) + ] + ) + + keyboard_rows.append( + [ + types.InlineKeyboardButton( + text="⬅️ Назад", + callback_data=f"admin_server_edit_{server.id}", + ) + ] + ) + + await callback.message.edit_text( + "\n".join(text_lines), + reply_markup=types.InlineKeyboardMarkup(inline_keyboard=keyboard_rows), + parse_mode="HTML", + ) + await callback.answer() + + +@admin_required +@error_handler +async def set_server_category( + callback: types.CallbackQuery, + db_user: User, + db: AsyncSession, +): + parts = callback.data.split("_") + server_id = int(parts[-2]) + category_id = int(parts[-1]) + + server = await update_server_squad(db, server_id, category_id=category_id) + + if not server: + await callback.answer("❌ Не удалось обновить категорию", show_alert=True) + return + + await cache.delete_pattern("available_countries*") + + text, keyboard = _build_server_edit_view(server) + await callback.message.edit_text( + text, + reply_markup=keyboard, + parse_mode="HTML", + ) + await callback.answer("✅ Категория обновлена") + + +@admin_required +@error_handler +async def clear_server_category( + callback: types.CallbackQuery, + db_user: User, + db: AsyncSession, +): + server_id = int(callback.data.split("_")[-1]) + + server = await update_server_squad(db, server_id, category_id=None) + + if not server: + await callback.answer("❌ Не удалось обновить категорию", show_alert=True) + return + + await cache.delete_pattern("available_countries*") + + text, keyboard = _build_server_edit_view(server) + await callback.message.edit_text( + text, + reply_markup=keyboard, + parse_mode="HTML", + ) + await callback.answer("✅ Категория удалена") + + +@admin_required +@error_handler +async def start_server_category_creation( + callback: types.CallbackQuery, + state: FSMContext, + db_user: User, + db: AsyncSession, +): + server_id = int(callback.data.split("_")[-1]) + server = await get_server_squad_by_id(db, server_id) + + if not server: + await callback.answer("❌ Сервер не найден!", show_alert=True) + return + + await state.set_data({"server_id": server_id}) + await state.set_state(AdminStates.creating_server_category) + + await callback.message.edit_text( + "🏷 Новая категория\n\n" + "Отправьте название новой категории для сервера:", + reply_markup=types.InlineKeyboardMarkup( + inline_keyboard=[ + [ + types.InlineKeyboardButton( + text="⬅️ Назад", callback_data=f"admin_server_edit_category_{server_id}" + ) + ] + ] + ), + parse_mode="HTML", + ) + await callback.answer() + + +@admin_required +@error_handler +async def process_server_category_creation( + message: types.Message, + state: FSMContext, + db_user: User, + db: AsyncSession, +): + data = await state.get_data() + server_id = data.get("server_id") + + if not server_id: + await message.answer("❌ Сервер не найден") + return + + category_name = (message.text or "").strip() + + if not (2 <= len(category_name) <= 255): + await message.answer("❌ Название категории должно содержать от 2 до 255 символов") + return + + try: + category = await create_server_category(db, name=category_name) + except Exception as error: + logger.error("Не удалось создать категорию %s: %s", category_name, error) + await message.answer("❌ Не удалось создать категорию. Попробуйте другое название") + return + + server = await update_server_squad(db, server_id, category_id=category.id) + + await state.clear() + + await cache.delete_pattern("available_countries*") + + if not server: + await message.answer("❌ Категория создана, но не удалось обновить сервер") + return + + safe_name = html.escape(category.name) + await message.answer( + f"✅ Категория {safe_name} создана и назначена серверу", + reply_markup=types.InlineKeyboardMarkup( + inline_keyboard=[ + [ + types.InlineKeyboardButton( + text="🔙 К серверу", callback_data=f"admin_server_edit_{server_id}" + ) + ] + ] + ), + parse_mode="HTML", + ) + + @admin_required @error_handler async def start_server_edit_promo_groups( @@ -1294,7 +1526,8 @@ def register_handlers(dp: Dispatcher): & ~F.data.contains("country") & ~F.data.contains("limit") & ~F.data.contains("desc") - & ~F.data.contains("promo"), + & ~F.data.contains("promo") + & ~F.data.contains("category"), ) dp.callback_query.register(toggle_server_availability, F.data.startswith("admin_server_toggle_")) dp.callback_query.register(toggle_server_trial_assignment, F.data.startswith("admin_server_trial_")) @@ -1304,14 +1537,19 @@ def register_handlers(dp: Dispatcher): dp.callback_query.register(start_server_edit_price, F.data.startswith("admin_server_edit_price_")) dp.callback_query.register(start_server_edit_country, F.data.startswith("admin_server_edit_country_")) dp.callback_query.register(start_server_edit_promo_groups, F.data.startswith("admin_server_edit_promo_")) - dp.callback_query.register(start_server_edit_limit, F.data.startswith("admin_server_edit_limit_")) - dp.callback_query.register(start_server_edit_description, F.data.startswith("admin_server_edit_desc_")) - + dp.callback_query.register(start_server_edit_limit, F.data.startswith("admin_server_edit_limit_")) + dp.callback_query.register(start_server_edit_description, F.data.startswith("admin_server_edit_desc_")) + dp.callback_query.register(show_server_category_menu, F.data.startswith("admin_server_edit_category_")) + dp.callback_query.register(set_server_category, F.data.startswith("admin_server_category_set_")) + dp.callback_query.register(clear_server_category, F.data.startswith("admin_server_category_clear_")) + dp.callback_query.register(start_server_category_creation, F.data.startswith("admin_server_category_create_")) + dp.message.register(process_server_name_edit, AdminStates.editing_server_name) dp.message.register(process_server_price_edit, AdminStates.editing_server_price) - dp.message.register(process_server_country_edit, AdminStates.editing_server_country) - dp.message.register(process_server_limit_edit, AdminStates.editing_server_limit) + dp.message.register(process_server_country_edit, AdminStates.editing_server_country) + dp.message.register(process_server_limit_edit, AdminStates.editing_server_limit) dp.message.register(process_server_description_edit, AdminStates.editing_server_description) + dp.message.register(process_server_category_creation, AdminStates.creating_server_category) dp.callback_query.register(toggle_server_promo_group, F.data.startswith("admin_server_promo_toggle_")) dp.callback_query.register(save_server_promo_groups, F.data.startswith("admin_server_promo_save_")) diff --git a/app/services/backup_service.py b/app/services/backup_service.py index abfde7e5..d62eec04 100644 --- a/app/services/backup_service.py +++ b/app/services/backup_service.py @@ -21,7 +21,7 @@ from app.database.models import ( User, Subscription, Transaction, PromoCode, PromoCodeUse, ReferralEarning, Squad, ServiceRule, SystemSetting, MonitoringLog, SubscriptionConversion, SentNotification, BroadcastHistory, - ServerSquad, SubscriptionServer, UserMessage, YooKassaPayment, + ServerCategory, ServerSquad, SubscriptionServer, UserMessage, YooKassaPayment, CryptoBotPayment, WelcomeText, Base, PromoGroup, AdvertisingCampaign, AdvertisingCampaignRegistration, SupportAuditLog, Ticket, TicketMessage, MulenPayPayment, Pal24Payment, DiscountOffer, WebApiToken, @@ -68,6 +68,7 @@ class BackupService: SystemSetting, ServiceRule, Squad, + ServerCategory, ServerSquad, PromoGroup, User, @@ -726,7 +727,7 @@ class BackupService: "mulenpay_payments", "pal24_payments", "transactions", "welcome_texts", "subscriptions", "promocodes", "users", "promo_groups", - "server_squads", "squads", "service_rules", + "server_squads", "server_categories", "squads", "service_rules", "system_settings", "web_api_tokens", "monitoring_logs" ] diff --git a/app/services/subscription_purchase_service.py b/app/services/subscription_purchase_service.py index 911c4d4c..a5f90a26 100644 --- a/app/services/subscription_purchase_service.py +++ b/app/services/subscription_purchase_service.py @@ -12,6 +12,7 @@ from app.database.crud.server_squad import ( get_available_server_squads, get_server_ids_by_uuids, get_server_squad_by_uuid, + choose_least_loaded_server_in_category, ) from app.database.crud.subscription import ( add_subscription_servers, @@ -103,6 +104,8 @@ class PurchaseServerOption: original_price_label: Optional[str] = None discount_percent: int = 0 is_available: bool = True + category_id: Optional[int] = None + is_category: bool = False def to_payload(self) -> Dict[str, Any]: payload: Dict[str, Any] = { @@ -119,6 +122,10 @@ class PurchaseServerOption: payload["original_price_label"] = self.original_price_label if self.discount_percent: payload["discount_percent"] = self.discount_percent + if self.category_id is not None: + payload["category_id"] = self.category_id + if self.is_category: + payload["is_category"] = True return payload @@ -129,6 +136,8 @@ class PurchaseServersConfig: max_selectable: int default_selection: List[str] hint: Optional[str] = None + category_mapping: Dict[str, List[str]] = field(default_factory=dict) + is_category_based: bool = False def to_payload(self) -> Dict[str, Any]: payload: Dict[str, Any] = { @@ -140,6 +149,10 @@ class PurchaseServersConfig: } if self.hint: payload["hint"] = self.hint + if self.category_mapping: + payload["category_mapping"] = self.category_mapping + if self.is_category_based: + payload["is_category_based"] = True return payload @@ -255,6 +268,7 @@ class PurchaseOptionsContext: default_period: PurchasePeriodConfig period_map: Dict[str, PurchasePeriodConfig] server_uuid_to_id: Dict[str, int] + category_mapping: Dict[str, List[str]] payload: Dict[str, Any] @@ -413,6 +427,7 @@ class MiniAppSubscriptionPurchaseService: user, texts, period_days, + available_servers, server_catalog, default_connected, ) @@ -483,6 +498,8 @@ class MiniAppSubscriptionPurchaseService: "selection": default_selection, "summary": None, } + if default_period.servers.category_mapping: + payload["server_category_mapping"] = default_period.servers.category_mapping return PurchaseOptionsContext( user=user, @@ -493,6 +510,7 @@ class MiniAppSubscriptionPurchaseService: default_period=default_period, period_map=period_map, server_uuid_to_id=server_uuid_to_id, + category_mapping=default_period.servers.category_mapping, payload=payload, ) @@ -566,25 +584,123 @@ class MiniAppSubscriptionPurchaseService: user: User, texts, period_days: int, + available_servers: List[ServerSquad], server_catalog: Dict[str, ServerSquad], default_selection: List[str], ) -> PurchaseServersConfig: discount_percent = user.get_promo_discount("servers", period_days) - options: List[PurchaseServerOption] = [] + default_selection_set = set(default_selection) + available_server_ids = {server.squad_uuid for server in available_servers} + + category_data: Dict[str, Dict[str, Any]] = {} + uncategorized_servers: List[ServerSquad] = [] for uuid, server in server_catalog.items(): + include_server = uuid in available_server_ids or uuid in default_selection_set + if not include_server: + continue + + if server.category_id and getattr(server, "category", None): + key = f"category:{server.category_id}" + data = category_data.setdefault( + key, + { + "category": server.category, + "servers": [], + "min_price": None, + "min_available_price": None, + "has_available": False, + }, + ) + data["servers"].append(server) + + price = server.price_kopeks or 0 + if data["min_price"] is None: + data["min_price"] = price + else: + data["min_price"] = min(data["min_price"], price) + + if getattr(server, "is_available", True) and not getattr(server, "is_full", False): + data["has_available"] = True + if data["min_available_price"] is None: + data["min_available_price"] = price + else: + data["min_available_price"] = min(data["min_available_price"], price) + else: + uncategorized_servers.append(server) + + category_items = list(category_data.items()) + category_items.sort( + key=lambda item: ( + getattr(item[1]["category"], "sort_order", 0), + getattr(item[1]["category"], "name", item[0]), + ) + ) + + options: List[PurchaseServerOption] = [] + category_mapping: Dict[str, List[str]] = {} + + for key, data in category_items: + category = data["category"] + base_price = data["min_available_price"] + if base_price is None: + base_price = data["min_price"] or 0 + discounted_per_month, discount_value = _apply_percentage_discount(base_price, discount_percent) + option = PurchaseServerOption( + uuid=key, + name=getattr(category, "name", key), + price_per_month=discounted_per_month, + price_label=texts.format_price(discounted_per_month), + original_price_per_month=base_price, + original_price_label=( + texts.format_price(base_price) + if discount_value and base_price != discounted_per_month + else None + ), + discount_percent=max(0, discount_percent), + is_available=data["has_available"], + category_id=getattr(category, "id", None), + is_category=True, + ) + options.append(option) + category_mapping[key] = [srv.squad_uuid for srv in data["servers"]] + + for server in uncategorized_servers: option = _build_server_option(server, discount_percent, texts) options.append(option) if not options: - default_selection = [] + default_values: List[str] = [] + else: + option_lookup = {option.uuid for option in options} + resolved_defaults: List[str] = [] + seen: set[str] = set() + + for uuid in default_selection: + server = server_catalog.get(uuid) + if server and server.category_id: + key = f"category:{server.category_id}" + if key in option_lookup and key not in seen: + resolved_defaults.append(key) + seen.add(key) + continue + if uuid in option_lookup and uuid not in seen: + resolved_defaults.append(uuid) + seen.add(uuid) + + if not resolved_defaults: + resolved_defaults = [options[0].uuid] + + default_values = resolved_defaults return PurchaseServersConfig( options=options, min_selectable=1 if options else 0, max_selectable=len(options), - default_selection=default_selection if default_selection else [opt.uuid for opt in options[:1]], + default_selection=default_values, hint=None, + category_mapping=category_mapping, + is_category_based=bool(category_mapping), ) def _build_devices_config( @@ -623,6 +739,40 @@ class MiniAppSubscriptionPurchaseService: hint=None, ) + async def _resolve_selected_servers( + self, + db: AsyncSession, + context: PurchaseOptionsContext, + selected_servers: List[str], + ) -> List[str]: + resolved: List[str] = [] + promo_group_id = getattr(context.user, "promo_group_id", None) + + for identifier in selected_servers: + if identifier.startswith("category:"): + if identifier not in context.category_mapping: + raise PurchaseValidationError("Invalid server category selection", code="invalid_servers") + try: + category_id = int(identifier.split(":", 1)[1]) + except (TypeError, ValueError) as error: + raise PurchaseValidationError("Invalid server category selection", code="invalid_servers") from error + + squad = await choose_least_loaded_server_in_category( + db, + category_id=category_id, + promo_group_id=promo_group_id, + ) + if not squad: + raise PurchaseValidationError( + "No available servers in selected category", + code="no_servers_in_category", + ) + resolved.append(squad.squad_uuid) + else: + resolved.append(identifier) + + return resolved + def parse_selection( self, context: PurchaseOptionsContext, @@ -722,14 +872,22 @@ class MiniAppSubscriptionPurchaseService: texts = get_texts(getattr(context.user, "language", None)) months = selection.period.months - server_ids = await get_server_ids_by_uuids(db, selection.servers) - if len(server_ids) != len(selection.servers): + resolved_servers = await self._resolve_selected_servers(db, context, selection.servers) + resolved_selection = PurchaseSelection( + period=selection.period, + traffic_value=selection.traffic_value, + servers=resolved_servers, + devices=selection.devices, + ) + + server_ids = await get_server_ids_by_uuids(db, resolved_selection.servers) + if len(server_ids) != len(resolved_selection.servers): raise PurchaseValidationError("Some selected servers are not available", code="invalid_servers") total_without_promo, details = await self._calculate_base_total( db, context.user, - selection, + resolved_selection, server_ids, ) @@ -768,7 +926,7 @@ class MiniAppSubscriptionPurchaseService: raise PurchaseValidationError("Failed to validate pricing", code="calculation_error") return PurchasePricingResult( - selection=selection, + selection=resolved_selection, server_ids=server_ids, server_prices_for_period=list(details.get("servers_individual_prices", [])), base_original_total=base_original_total, diff --git a/app/states.py b/app/states.py index 39b3744d..562a7934 100644 --- a/app/states.py +++ b/app/states.py @@ -113,6 +113,7 @@ class AdminStates(StatesGroup): editing_server_limit = State() editing_server_description = State() editing_server_promo_groups = State() + creating_server_category = State() creating_server_uuid = State() creating_server_name = State() diff --git a/tests/services/test_subscription_auto_purchase_service.py b/tests/services/test_subscription_auto_purchase_service.py index 123b0cdc..186944ee 100644 --- a/tests/services/test_subscription_auto_purchase_service.py +++ b/tests/services/test_subscription_auto_purchase_service.py @@ -1,11 +1,14 @@ import pytest from datetime import datetime, timedelta +from typing import Any, Dict from unittest.mock import AsyncMock, MagicMock from app.config import settings from app.database.models import User +from app.database.crud import server_squad as server_squad_crud from app.services.subscription_auto_purchase_service import auto_purchase_saved_cart_after_topup from app.services.subscription_purchase_service import ( + MiniAppSubscriptionPurchaseService, PurchaseDevicesConfig, PurchaseOptionsContext, PurchasePeriodConfig, @@ -92,6 +95,7 @@ async def test_auto_purchase_saved_cart_after_topup_success(monkeypatch): default_period=period_config, period_map={"days:30": period_config}, server_uuid_to_id={"ru": 1}, + category_mapping={}, payload={}, ) @@ -701,3 +705,110 @@ async def test_auto_purchase_trial_remaining_days_transferred(monkeypatch): expected_end = trial_end + timedelta(days=32) # trial_end + (30 + 2) actual_delta = (subscription.end_date - trial_end).days assert actual_delta == 32, f"Expected 32 days extension (30 + 2 bonus), got {actual_delta}" + + +def test_build_servers_config_category_price_ignores_full_squads(): + service = MiniAppSubscriptionPurchaseService() + + user = MagicMock(spec=User) + user.get_promo_discount = MagicMock(return_value=0) + + texts = DummyTexts() + + category = MagicMock() + category.id = 7 + category.name = "Category" + category.sort_order = 1 + + full_server = MagicMock() + full_server.squad_uuid = "srv-cheap" + full_server.category_id = category.id + full_server.category = category + full_server.price_kopeks = 10_000 + full_server.is_available = True + full_server.is_full = True + + available_server = MagicMock() + available_server.squad_uuid = "srv-available" + available_server.category_id = category.id + available_server.category = category + available_server.price_kopeks = 20_000 + available_server.is_available = True + available_server.is_full = False + + available_servers = [full_server, available_server] + server_catalog = {srv.squad_uuid: srv for srv in available_servers} + + config = service._build_servers_config( + user, + texts, + period_days=30, + available_servers=available_servers, + server_catalog=server_catalog, + default_selection=[], + ) + + category_option = next(option for option in config.options if option.is_category) + + assert category_option.price_per_month == available_server.price_kopeks + assert category_option.original_price_per_month == available_server.price_kopeks + assert config.category_mapping[f"category:{category.id}"] == [ + full_server.squad_uuid, + available_server.squad_uuid, + ] + + +@pytest.mark.asyncio +async def test_choose_least_loaded_server_in_category_prefers_lowest_load(monkeypatch): + capture: Dict[str, Any] = {} + + async def fake_get_available(db, promo_group_id=None, category_id=None, only_with_capacity=False): + capture["only_with_capacity"] = only_with_capacity + high_load = MagicMock() + high_load.max_users = 100 + high_load.current_users = 80 + high_load.id = 10 + + low_load = MagicMock() + low_load.max_users = 200 + low_load.current_users = 20 + low_load.id = 5 + return [high_load, low_load] + + monkeypatch.setattr( + server_squad_crud, + "get_available_server_squads", + fake_get_available, + ) + + dummy_db = AsyncMock(spec=AsyncSession) + + squad = await server_squad_crud.choose_least_loaded_server_in_category( + dummy_db, + category_id=1, + promo_group_id=2, + ) + + assert squad.current_users == 20 + assert capture["only_with_capacity"] is True + + +@pytest.mark.asyncio +async def test_choose_least_loaded_server_in_category_returns_none_when_empty(monkeypatch): + async def fake_get_available(db, promo_group_id=None, category_id=None, only_with_capacity=False): + return [] + + monkeypatch.setattr( + server_squad_crud, + "get_available_server_squads", + fake_get_available, + ) + + dummy_db = AsyncMock(spec=AsyncSession) + + squad = await server_squad_crud.choose_least_loaded_server_in_category( + dummy_db, + category_id=5, + ) + + assert squad is None