197 lines
9.0 KiB
Python
197 lines
9.0 KiB
Python
#!/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()) |