#!/usr/bin/env python3 """ Скрипт для очистки и инициализации БД с тестовыми данными """ import asyncio import random from database import async_session_maker, Base, engine from models import User, Lottery, Participation, Winner from services import LotteryService, UserService, ParticipationService from datetime import datetime, timedelta from sqlalchemy import delete, update, select, func async def clear_database(): """Очистка всех таблиц""" print("🗑️ Очистка базы данных...") async with async_session_maker() as session: # Удаляем все данные в правильном порядке (учитывая внешние ключи) await session.execute(delete(Winner)) await session.execute(delete(Participation)) await session.execute(delete(Lottery)) await session.execute(delete(User)) await session.commit() print("✅ База данных очищена") def generate_account_numbers(user_id: int, count: int = 10) -> list: """Генерация номеров счетов для пользователя""" base = user_id * 1000 return [f"{base + i:06d}" for i in range(1, count + 1)] async def create_test_users(): """Создание тестовых пользователей""" print("👥 Создание тестовых пользователей...") users_data = [ {"telegram_id": 111111111, "username": "alice_winner", "first_name": "Alice", "last_name": "Johnson"}, {"telegram_id": 222222222, "username": "bob_player", "first_name": "Bob", "last_name": "Smith"}, {"telegram_id": 333333333, "username": "charlie_test", "first_name": "Charlie", "last_name": "Brown"}, {"telegram_id": 444444444, "username": "diana_luck", "first_name": "Diana", "last_name": "Wilson"}, {"telegram_id": 555555555, "username": "eve_gamer", "first_name": "Eve", "last_name": "Davis"}, {"telegram_id": 666666666, "username": "frank_player", "first_name": "Frank", "last_name": "Miller"}, {"telegram_id": 777777777, "username": "grace_lucky", "first_name": "Grace", "last_name": "Taylor"}, {"telegram_id": 888888888, "username": "henry_bet", "first_name": "Henry", "last_name": "Anderson"}, {"telegram_id": 999999999, "username": "iris_game", "first_name": "Iris", "last_name": "Thomas"}, {"telegram_id": 101010101, "username": "jack_chance", "first_name": "Jack", "last_name": "Jackson"}, ] created_users = [] async with async_session_maker() as session: for i, user_data in enumerate(users_data, 1): # Создаем пользователя user = await UserService.get_or_create_user( session=session, telegram_id=user_data["telegram_id"], username=user_data["username"], first_name=user_data["first_name"], last_name=user_data.get("last_name", None) ) # Генерируем и присваиваем номера счетов account_numbers = generate_account_numbers(i, 10) # Обновляем счета await session.execute( update(User) .where(User.id == user.id) .values(account_number=",".join(account_numbers)) ) await session.commit() created_users.append((user, account_numbers)) print(f"✅ Создан пользователь {user.first_name} ({user.username}) с {len(account_numbers)} счетами") return created_users async def create_test_lotteries(users_data): """Создание тестовых розыгрышей""" print("🎲 Создание тестовых розыгрышей...") async with async_session_maker() as session: # Первый розыгрыш - активный lottery1 = await LotteryService.create_lottery( session=session, title="🎯 Супер Розыгрыш #1", description="Розыгрыш крутых призов! Участвуют все желающие.", prizes=[ "iPhone 15 Pro - Новейший флагман Apple", "AirPods Pro - Беспроводные наушники премиум класса", "Apple Watch - Умные часы последней модели" ], creator_id=users_data[0][0].id ) # Второй розыгрыш - тоже активный lottery2 = await LotteryService.create_lottery( session=session, title="🚀 Мега Розыгрыш #2", description="Еще больше призов! Еще больше возможностей!", prizes=[ "MacBook Pro - Мощный ноутбук для профессионалов", "iPad Air - Планшет для творчества и работы", "Денежный приз 50,000 рублей" ], creator_id=users_data[0][0].id ) await session.commit() print(f"✅ Создан розыгрыш: {lottery1.title} (ID: {lottery1.id})") print(f"✅ Создан розыгрыш: {lottery2.title} (ID: {lottery2.id})") return [lottery1, lottery2] async def populate_participants(users_data, lotteries): """Заполнение участников розыгрышей""" print("🎫 Добавление участников в розыгрыши...") async with async_session_maker() as session: # Добавляем всех пользователей в первый розыгрыш lottery1 = lotteries[0] for user, account_numbers in users_data: # Добавляем каждого пользователя как участника success = await ParticipationService.add_participant( session=session, lottery_id=lottery1.id, user_id=user.id ) if success: print(f"✅ Пользователь {user.first_name}: добавлен в розыгрыш №1") else: print(f"⚠️ Пользователь {user.first_name}: уже участвует в розыгрыше №1") # Добавляем первых 5 пользователей во второй розыгрыш lottery2 = lotteries[1] for user, account_numbers in users_data[:5]: # Первые 5 пользователей success = await ParticipationService.add_participant( session=session, lottery_id=lottery2.id, user_id=user.id ) if success: print(f"✅ Пользователь {user.first_name}: добавлен в розыгрыш №2") else: print(f"⚠️ Пользователь {user.first_name}: уже участвует в розыгрыше №2") async def show_statistics(): """Показ статистики созданных данных""" print("\n📊 Статистика созданных данных:") async with async_session_maker() as session: # Пользователи users = await UserService.get_all_users(session) print(f"👥 Пользователей: {len(users)}") # Розыгрыши lotteries = await LotteryService.get_active_lotteries(session) print(f"🎲 Активных розыгрышей: {len(lotteries)}") # Участия for lottery in lotteries: # Простой подсчет через SQL result = await session.execute( select(func.count(Participation.id)) .where(Participation.lottery_id == lottery.id) ) participation_count = result.scalar() print(f"🎫 Розыгрыш '{lottery.title}': {participation_count} участий") print("\n✅ Тестовые данные успешно созданы!") async def main(): """Главная функция""" print("🚀 Инициализация тестовой базы данных") print("=" * 50) # Очищаем БД await clear_database() # Создаем тестовых пользователей users_data = await create_test_users() # Создаем розыгрыши lotteries = await create_test_lotteries(users_data) # Заполняем участников await populate_participants(users_data, lotteries) # Показываем статистику await show_statistics() print("\n🎉 Инициализация завершена!") if __name__ == "__main__": asyncio.run(main())