Files
New-planet-ai-agent/scripts/generate_test_data.py

116 lines
4.4 KiB
Python
Raw Permalink Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
"""Скрипт для генерации тестовых данных."""
import argparse
import asyncio
import json
import random
from datetime import date, timedelta
from uuid import uuid4
from models.schedule import Schedule, Task
def generate_tasks(count: int = 5) -> list[Task]:
"""Генерировать тестовые задания."""
task_templates = [
{"title": "Утренняя зарядка", "category": "утренняя_рутина", "duration": 15},
{"title": "Чистка зубов", "category": "утренняя_рутина", "duration": 5},
{"title": "Завтрак", "category": "утренняя_рутина", "duration": 20},
{"title": "Рисование", "category": "обучение", "duration": 30},
{"title": "Чтение книги", "category": "обучение", "duration": 20},
{"title": "Игра с конструктором", "category": "игра", "duration": 45},
{"title": "Прогулка", "category": "игра", "duration": 60},
{"title": "Обед", "category": "отдых", "duration": 30},
{"title": "Тихий час", "category": "отдых", "duration": 60},
{"title": "Ужин", "category": "вечерняя_рутина", "duration": 30},
{"title": "Подготовка ко сну", "category": "вечерняя_рутина", "duration": 20},
]
selected = random.sample(task_templates, min(count, len(task_templates)))
tasks = []
for i, template in enumerate(selected):
tasks.append(
Task(
id=uuid4(),
title=template["title"],
description=f"Описание для {template['title']}",
duration_minutes=template["duration"],
category=template["category"],
completed=random.choice([True, False]),
order=i,
)
)
return tasks
def generate_schedules(user_id: str, count: int, start_date: date = None) -> list[Schedule]:
"""Генерировать тестовые расписания."""
if start_date is None:
start_date = date.today()
schedules = []
for i in range(count):
schedule_date = start_date + timedelta(days=i)
tasks = generate_tasks(random.randint(4, 8))
schedules.append(
Schedule(
id=uuid4(),
title=f"Расписание на {schedule_date.strftime('%d.%m.%Y')}",
date=schedule_date,
tasks=tasks,
user_id=user_id,
created_at=schedule_date.isoformat(),
)
)
return schedules
async def main():
"""Главная функция."""
parser = argparse.ArgumentParser(description="Генерация тестовых данных")
parser.add_argument("--users", type=int, default=10, help="Количество пользователей")
parser.add_argument("--schedules", type=int, default=50, help="Количество расписаний")
parser.add_argument("--output", type=str, default="test_data.json", help="Файл для сохранения")
args = parser.parse_args()
print(f"Генерация тестовых данных:")
print(f" Пользователей: {args.users}")
print(f" Расписаний: {args.schedules}")
data = {
"users": [],
"schedules": [],
}
# Генерируем пользователей
for _ in range(args.users):
user_id = str(uuid4())
data["users"].append(
{
"id": user_id,
"email": f"user_{random.randint(1000, 9999)}@example.com",
}
)
# Генерируем расписания для пользователя
schedules_per_user = args.schedules // args.users
user_schedules = generate_schedules(user_id, schedules_per_user)
data["schedules"].extend([s.model_dump() for s in user_schedules])
# Сохраняем в файл
with open(args.output, "w", encoding="utf-8") as f:
json.dump(data, f, ensure_ascii=False, indent=2, default=str)
print(f"\nДанные сохранены в {args.output}")
print(f" Пользователей: {len(data['users'])}")
print(f" Расписаний: {len(data['schedules'])}")
if __name__ == "__main__":
asyncio.run(main())