""" 创建测试数据脚本 """ 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())