buff/backend/init_test_data.py

"""
初始化测试数据
"""
from sqlalchemy.orm import Session
from app.database import SessionLocal
from app.models.user import User
from app.models.card import Card
from app.models.software import Software
from app.models.announcement import Announcement
import bcrypt
from datetime import datetime, timedelta
import random
import string

def generate_random_string(length=8):
    """生成随机字符串"""
    return ''.join(random.choices(string.ascii_letters + string.digits, k=length))

def init_test_data():
    """初始化测试数据"""
    db = SessionLocal()
    try:
        # 检查是否已有用户数据
        user_count = db.query(User).count()
        if user_count > 0:
            print(f"数据库中已有 {user_count} 个用户，跳过初始化")
            return
        
        # 获取软件列表
        software_list = db.query(Software).all()
        if not software_list:
            print("警告：没有找到软件数据，请先创建软件")
            return
        
        software = software_list[0]
        print(f"使用软件: {software.name} ({software.code})")
        
        # 创建测试用户
        print("正在创建测试用户...")
        test_users = []
        for i in range(1, 21):  # 创建20个测试用户
            username = f"testuser{i:02d}"
            email = f"test{i}@example.com"
            machine_code = f"MACHINE{generate_random_string(8).upper()}"
            
            # 随机设置过期时间
            if i % 3 == 0:
                expire_time = None  # 无有效期
            elif i % 3 == 1:
                expire_time = datetime.now() + timedelta(days=30)  # 有效
            else:
                expire_time = datetime.now() - timedelta(days=1)  # 已过期
            
            # 随机状态
            status = 1 if i % 5 != 0 else 0  # 大部分正常，少数禁用
            
            # 生成密码哈希
            password_hash = bcrypt.hashpw("123456".encode('utf-8'), bcrypt.gensalt()).decode('utf-8')
            
            user = User(
                username=username,
                password=password_hash,  # 默认密码
                email=email,
                machine_code=machine_code,
                expire_time=expire_time,
                status=status,
                created_at=datetime.now() - timedelta(days=random.randint(1, 90)),
                updated_at=datetime.now()
            )
            db.add(user)
            test_users.append(user)
        
        db.flush()  # 获取用户ID
        
        # 创建测试卡密
        print("正在创建测试卡密...")
        for i in range(1, 11):  # 创建10个测试卡密
            card_code = f"CARD{generate_random_string(10).upper()}"
            days = random.choice([7, 30, 90, 180, 365])
            
            # 随机分配给用户
            user = random.choice(test_users) if test_users else None
            
            card = Card(
                card_code=card_code,
                software_code=software.code,
                days=days,
                discount=random.choice([100, 95, 90, 85]),  # 折扣
                user_id=user.id if user else None,
                activated_at=datetime.now() - timedelta(days=random.randint(1, 30)) if user else None,
                expire_time=datetime.now() + timedelta(days=days) if user else None,
                status=1 if user else 0,  # 已激活或未激活
                created_at=datetime.now() - timedelta(days=random.randint(1, 60)),
                updated_at=datetime.now()
            )
            db.add(card)
            
            # 如果卡密已激活，更新用户的卡密ID和过期时间
            if user:
                user.card_id = card.id
                user.expire_time = card.expire_time
        
        # 创建测试公告
        print("正在创建测试公告...")
        announcements = [
            {
                "title": "系统维护通知",
                "content": "系统将于本周六晚上进行维护，预计维护时间2小时，期间可能无法正常使用，请提前做好准备。",
                "priority": 1,
                "is_active": 1
            },
            {
                "title": "新功能上线",
                "content": "我们很高兴地宣布，新版本已经上线，包含以下新功能：\n1. 优化了用户体验\n2. 增加了新的功能模块\n3. 修复了已知问题",
                "priority": 2,
                "is_active": 1
            },
            {
                "title": "优惠活动",
                "content": "限时优惠活动正在进行中，购买年卡可享受8折优惠，活动时间有限，请抓紧时间！",
                "priority": 3,
                "is_active": 1
            }
        ]
        
        for ann_data in announcements:
            announcement = Announcement(
                title=ann_data["title"],
                content=ann_data["content"],
                priority=ann_data["priority"],
                is_active=ann_data["is_active"],
                publish_time=datetime.now() - timedelta(days=random.randint(1, 30)),
                created_at=datetime.now() - timedelta(days=random.randint(1, 30)),
                updated_at=datetime.now()
            )
            db.add(announcement)
        
        db.commit()
        print("测试数据初始化完成！")
        print(f"- 创建了 {len(test_users)} 个测试用户")
        print(f"- 创建了 10 个测试卡密")
        print(f"- 创建了 {len(announcements)} 条测试公告")
        print("\n测试用户账号信息：")
        print("用户名格式: testuser01, testuser02, ...")
        print("默认密码: 123456")
        
    except Exception as e:
        db.rollback()
        print(f"初始化失败: {e}")
        import traceback
        traceback.print_exc()
        raise
    finally:
        db.close()

if __name__ == "__main__":
    init_test_data()