012-kaopeilian/backend/scripts/create_test_data.py

"""
创建测试数据脚本
"""
import asyncio
import sys
from pathlib import Path

# 添加项目路径
sys.path.append(str(Path(__file__).resolve().parent.parent))

from sqlalchemy.ext.asyncio import AsyncSession

from app.core.database import AsyncSessionLocal
from app.core.logger import logger
from app.core.security import get_password_hash
from app.models.user import Team, User
from app.schemas.user import UserCreate
from app.services.user_service import UserService


async def create_test_users(db: AsyncSession):
    """创建测试用户"""
    user_service = UserService(db)
    
    # 测试用户数据
    test_users = [
        {
            "username": "admin",
            "email": "admin@kaopeilian.com",
            "password": "admin123",
            "full_name": "系统管理员",
            "role": "admin",
        },
        {
            "username": "trainee1",
            "email": "trainee1@kaopeilian.com",
            "password": "trainee123",
            "full_name": "张三",
            "phone": "13800138001",
            "role": "trainee",
        },
        {
            "username": "trainee2",
            "email": "trainee2@kaopeilian.com",
            "password": "trainee123",
            "full_name": "李四",
            "phone": "13800138002",
            "role": "trainee",
        },
    ]
    
    created_users = []
    for user_data in test_users:
        # 检查用户是否已存在
        existing_user = await user_service.get_by_username(user_data["username"])
        if existing_user:
            logger.info(f"用户 {user_data['username']} 已存在，跳过创建")
            created_users.append(existing_user)
        else:
            # 创建用户
            user_create = UserCreate(**user_data)
            user = await user_service.create_user(obj_in=user_create)
            created_users.append(user)
            logger.info(f"创建用户: {user.username} ({user.role})")
    
    return created_users


async def create_test_teams(db: AsyncSession, users: list[User]):
    """创建测试团队"""
    # 获取管理员
    admin = next(u for u in users if u.role == "admin")
    
    # 检查团队是否已存在
    from sqlalchemy import select
    result = await db.execute(
        select(Team).where(Team.code == "TECH")
    )
    existing_team = result.scalar_one_or_none()
    
    if not existing_team:
        # 创建技术部
        tech_team = Team(
            name="技术部",
            code="TECH",
            description="负责产品研发和技术支持",
            team_type="department",
            leader_id=admin.id,
            created_by=admin.id,
        )
        db.add(tech_team)
        await db.commit()
        await db.refresh(tech_team)
        logger.info(f"创建团队: {tech_team.name}")
        
        # 创建前端组
        frontend_team = Team(
            name="前端开发组",
            code="TECH-FE",
            description="负责前端开发",
            team_type="project",
            parent_id=tech_team.id,
            created_by=admin.id,
        )
        db.add(frontend_team)
        
        # 创建后端组
        backend_team = Team(
            name="后端开发组",
            code="TECH-BE",
            description="负责后端开发",
            team_type="project",
            parent_id=tech_team.id,
            created_by=admin.id,
        )
        db.add(backend_team)
        
        await db.commit()
        logger.info("创建子团队: 前端开发组、后端开发组")
        
        # 将用户加入团队
        user_service = UserService(db)
        
        # 学员加入子团队
        for user in users:
            if user.role == "trainee":
                if "1" in user.username:
                    await user_service.add_user_to_team(
                        user_id=user.id,
                        team_id=frontend_team.id,
                        role="member"
                    )
                else:
                    await user_service.add_user_to_team(
                        user_id=user.id,
                        team_id=backend_team.id,
                        role="member"
                    )
        
        logger.info("用户已加入相应团队")
    else:
        logger.info("团队已存在，跳过创建")


async def main():
    """主函数"""
    async with AsyncSessionLocal() as db:
        try:
            logger.info("开始创建测试数据...")
            
            # 创建测试用户
            users = await create_test_users(db)
            
            # 创建测试团队
            await create_test_teams(db, users)
            
            logger.info("测试数据创建完成！")
            logger.info("\n可用的测试账号：")
            logger.info("管理员 - 用户名: admin, 密码: admin123")
            logger.info("学员 - 用户名: trainee1, 密码: trainee123")
            logger.info("学员 - 用户名: trainee2, 密码: trainee123")
            
        except Exception as e:
            logger.error(f"创建测试数据失败: {str(e)}")
            await db.rollback()
            raise


if __name__ == "__main__":
    asyncio.run(main())