Maas-Model-platform/backend/scripts/batch_create_users.py

"""
批量用户注册脚本

用于JMeter性能测试的用户数据准备，支持配置用户数量和并发批次

使用说明:
    python batch_create_users.py --count 1000 --batch-size 100 --output users.csv
    
参数说明:
    --count: 要创建的用户总数，默认1000
    --batch-size: 每批次插入的用户数，默认100
    --prefix: 用户名前缀，默认"testuser"
    --output: 输出文件路径，默认"test_users.csv"
    --password: 用户统一密码，默认"Test@123"
"""
import sys
import uuid
import csv
import argparse
from datetime import date
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent.parent))

from app.database import SessionLocal
from app.models.user import User
from app.services.auth_service import AuthService


def generate_users(count: int, prefix: str = "testuser", password: str = "Test@123") -> list:
    """
    生成用户数据列表
    
    Args:
        count: 用户数量
        prefix: 用户名前缀
        password: 用户密码
        
    Returns:
        list: 用户数据字典列表
    """
    users = []
    for i in range(count):
        user_id = str(uuid.uuid4())
        username = f"{prefix}{i:06d}"
        users.append({
            "id": user_id,
            "username": username,
            "password_hash": AuthService.hash_password(password),
            "nickname": f"测试用户{i + 1}",
            "email": f"{username}@test.com",
            "apikey": f"sk-test-{user_id[:8]}",
            "registration_date": date.today(),
            "status": "active"
        })
    return users


def batch_insert_users(db, users: list, batch_size: int = 100) -> int:
    """
    批量插入用户数据
    
    Args:
        db: 数据库会话
        users: 用户数据列表
        batch_size: 每批次插入数量
        
    Returns:
        int: 成功插入的用户数
    """
    total_inserted = 0
    total_batches = (len(users) + batch_size - 1) // batch_size
    
    for i in range(0, len(users), batch_size):
        batch = users[i:i + batch_size]
        db.bulk_insert_mappings(User, batch)
        db.commit()
        total_inserted += len(batch)
        
        if (i // batch_size + 1) % 10 == 0 or (i // batch_size + 1) == total_batches:
            print(f"已插入 {total_inserted}/{len(users)} 用户...")
    
    return total_inserted


def write_output_file(users: list, output_path: str, password: str) -> None:
    """
    将用户信息写入CSV文件
    
    Args:
        users: 用户数据列表
        output_path: 输出文件路径
        password: 用户密码
    """
    with open(output_path, 'w', newline='', encoding='utf-8') as f:
        writer = csv.writer(f)
        writer.writerow(['username', 'password', 'email', 'apikey', 'user_id'])
        for user in users:
            writer.writerow([
                user['username'],
                password,
                user['email'],
                user['apikey'],
                user['id']
            ])
    print(f"用户信息已写入: {output_path}")


def main():
    parser = argparse.ArgumentParser(description="批量创建测试用户")
    parser.add_argument("--count", type=int, default=1000, help="要创建的用户总数")
    parser.add_argument("--batch-size", type=int, default=100, help="每批次插入的用户数")
    parser.add_argument("--prefix", type=str, default="testuser", help="用户名前缀")
    parser.add_argument("--output", type=str, default="test_users.csv", help="输出文件路径")
    parser.add_argument("--password", type=str, default="Test@123", help="用户统一密码")
    
    args = parser.parse_args()
    
    print(f"开始批量创建用户...")
    print(f"  用户总数: {args.count}")
    print(f"  批次大小: {args.batch_size}")
    print(f"  用户名前缀: {args.prefix}")
    print(f"  输出文件: {args.output}")
    
    db = SessionLocal()
    
    try:
        print("\n1. 生成用户数据...")
        users = generate_users(args.count, args.prefix, args.password)
        print(f"   已生成 {len(users)} 个用户数据")
        
        print("\n2. 批量插入数据库...")
        inserted = batch_insert_users(db, users, args.batch_size)
        print(f"   成功插入 {inserted} 个用户")
        
        print("\n3. 写入输出文件...")
        write_output_file(users, args.output, args.password)
        
        print("\n✅ 批量用户创建完成!")
        print(f"   建议JMeter并发数: {min(args.count, 200)} (可根据服务器配置调整)")
        print(f"   如需更高并发，建议分批创建，每批不超过5000用户")
        
    except Exception as e:
        print(f"\n❌ 执行失败: {e}")
        db.rollback()
        raise
    finally:
        db.close()


if __name__ == "__main__":
    main()