generate_test_data.py

#!/usr/bin/env python
"""
生成测试数据脚本
为HR管理系统生成100条测试数据
"""

import os
import sys
import django
from datetime import date, datetime, timedelta
from decimal import Decimal
import random

# 设置Django环境
os.environ.setdefault('DJANGO_SETTINGS_MODULE', 'hr_system.settings')
django.setup()

from django.contrib.auth.models import User
from employees.models import Employee
from departments.models import Department
from salary.models import SalaryRecord, SalaryStructure, SalaryPayment
from recruitment.models import JobPosition, Candidate, JobApplication, Interview, OfferLetter
from attendance.models import AttendanceRecord, LeaveType, LeaveApplication, OvertimeRecord

def create_test_data():
    print("开始生成测试数据...")
    
    # 1. 创建更多部门
    departments_data = [
        {'name': '人工智能部', 'code': 'AI', 'description': '负责AI技术研发'},
        {'name': '数据科学部', 'code': 'DS', 'description': '负责数据分析和挖掘'},
        {'name': '云计算部', 'code': 'CLOUD', 'description': '负责云平台建设'},
        {'name': '移动开发部', 'code': 'MOBILE', 'description': '负责移动应用开发'},
        {'name': '测试部', 'code': 'QA', 'description': '负责产品质量保证'},
        {'name': '运维部', 'code': 'OPS', 'description': '负责系统运维'},
        {'name': '法务部', 'code': 'LEGAL', 'description': '负责法律事务'},
        {'name': '采购部', 'code': 'PURCHASE', 'description': '负责采购管理'},
    ]
    
    created_departments = []
    for dept_data in departments_data:
        dept, created = Department.objects.get_or_create(
            code=dept_data['code'],
            defaults={
                'name': dept_data['name'],
                'description': dept_data['description'],
                'is_active': True
            }
        )
        if created:
            created_departments.append(dept)
            print(f"创建部门: {dept.name}")
    
    # 2. 创建更多员工
    all_departments = list(Department.objects.filter(is_active=True))
    
    # 员工姓名列表
    first_names = ['张', '李', '王', '刘', '陈', '杨', '赵', '黄', '周', '吴', '徐', '孙', '胡', '朱', '高', '林', '何', '郭', '马', '罗']
    last_names = ['伟', '芳', '娜', '秀英', '敏', '静', '丽', '强', '磊', '军', '洋', '勇', '艳', '杰', '娟', '涛', '明', '超', '秀兰', '霞']
    
    positions = [
        '软件工程师', '高级工程师', '架构师', '产品经理', '项目经理', 
        '测试工程师', '运维工程师', '数据分析师', 'UI设计师', '前端工程师',
        '后端工程师', '全栈工程师', '算法工程师', '机器学习工程师', '数据科学家',
        '技术总监', '研发经理', '质量经理', '运营专员', '市场专员'
    ]
    
    employees_created = 0
    for i in range(50):  # 创建50个员工
        first_name = random.choice(first_names)
        last_name = random.choice(last_names)
        name = first_name + last_name
        
        # 检查是否已存在
        employee_id = f"EMP{2024}{str(i+100).zfill(3)}"
        if Employee.objects.filter(employee_id=employee_id).exists():
            continue
            
        try:
            # 创建用户
            username = f"user{i+100}"
            if User.objects.filter(username=username).exists():
                continue
                
            user = User.objects.create_user(
                username=username,
                email=f"{username}@company.com",
                password="123456",
                first_name=first_name,
                last_name=last_name
            )
            
            # 创建员工
            employee = Employee.objects.create(
                user=user,
                employee_id=employee_id,
                name=name,
                department=random.choice(all_departments),
                position=random.choice(positions),
                phone=f"138{random.randint(10000000, 99999999)}",
                email=f"{username}@company.com",
                id_card=f"{random.randint(100000, 999999)}{random.randint(1985, 2000)}{str(random.randint(1, 12)).zfill(2)}{str(random.randint(1, 28)).zfill(2)}{random.randint(1000, 9999)}",
                hire_date=date.today() - timedelta(days=random.randint(30, 1000)),
                salary=Decimal(str(random.randint(8000, 30000))),
                gender=random.choice(['M', 'F']),
                birth_date=date(random.randint(1985, 2000), random.randint(1, 12), random.randint(1, 28)),
                is_active=True
            )
            employees_created += 1
            if employees_created % 10 == 0:
                print(f"已创建 {employees_created} 个员工...")
                
        except Exception as e:
            print(f"创建员工失败: {e}")
            continue
    
    print(f"总共创建了 {employees_created} 个员工")
    
    # 3. 创建薪资记录
    all_employees = list(Employee.objects.filter(is_active=True))
    all_structures = list(SalaryStructure.objects.filter(is_active=True))
    
    if all_structures:
        salary_records_created = 0
        current_year = date.today().year
        
        for employee in all_employees[:30]:  # 为前30个员工创建薪资记录
            for month in range(1, 7):  # 创建前6个月的记录
                try:
                    salary_record = SalaryRecord.objects.create(
                        employee=employee,
                        year=current_year,
                        month=month,
                        salary_structure=random.choice(all_structures),
                        basic_salary=employee.salary or Decimal('10000'),
                        housing_allowance=Decimal(str(random.randint(500, 2000))),
                        transport_allowance=Decimal(str(random.randint(200, 800))),
                        meal_allowance=Decimal(str(random.randint(300, 600))),
                        communication_allowance=Decimal(str(random.randint(100, 300))),
                        performance_bonus=Decimal(str(random.randint(0, 5000))),
                        overtime_pay=Decimal(str(random.randint(0, 2000))),
                        other_bonus=Decimal(str(random.randint(0, 1000))),
                        social_insurance=Decimal(str(random.randint(800, 1500))),
                        housing_fund=Decimal(str(random.randint(600, 1200))),
                        other_deduction=Decimal(str(random.randint(0, 500))),
                        work_days=22,
                        actual_work_days=random.randint(20, 22),
                        overtime_hours=Decimal(str(random.randint(0, 40))),
                        status=random.choice(['calculated', 'approved', 'paid']),
                        remark=f"{current_year}年{month}月薪资"
                    )
                    # 计算薪资
                    salary_record.calculate_salary()
                    salary_record.save()
                    salary_records_created += 1
                    
                except Exception as e:
                    print(f"创建薪资记录失败: {e}")
                    continue
        
        print(f"创建了 {salary_records_created} 条薪资记录")
    
    # 4. 创建招聘数据
    # 创建更多职位
    job_titles = [
        'Python开发工程师', 'Java开发工程师', 'Go开发工程师', 'React前端工程师',
        '产品经理', '数据分析师', 'AI算法工程师', '测试工程师', '运维工程师',
        'UI/UX设计师', '项目经理', '技术总监', '架构师', '全栈工程师'
    ]
    
    positions_created = 0
    for title in job_titles:
        try:
            position = JobPosition.objects.create(
                title=title,
                department=random.choice(all_departments),
                description=f"{title}职位描述，负责相关技术开发和维护工作。",
                requirements=f"要求具备{title}相关技能，有相关工作经验优先。",
                salary_min=Decimal(str(random.randint(8000, 15000))),
                salary_max=Decimal(str(random.randint(15000, 30000))),
                location='北京',
                employment_type='full_time',
                experience_required=random.choice(['entry', 'junior', 'mid', 'senior']),
                education_required=random.choice(['bachelor', 'master']),
                status='published',
                recruiter=User.objects.filter(is_staff=True).first()
            )
            positions_created += 1
        except Exception as e:
            print(f"创建职位失败: {e}")
            continue
    
    print(f"创建了 {positions_created} 个职位")
    
    # 创建候选人
    candidates_created = 0
    for i in range(30):
        first_name = random.choice(first_names)
        last_name = random.choice(last_names)
        name = first_name + last_name
        
        try:
            candidate = Candidate.objects.create(
                name=name,
                email=f"candidate{i+100}@email.com",
                phone=f"139{random.randint(10000000, 99999999)}",
                gender=random.choice(['M', 'F']),
                education=random.choice(['bachelor', 'master', 'doctor']),
                current_company=f"公司{i+1}",
                current_position=random.choice(positions),
                expected_salary=Decimal(str(random.randint(10000, 25000)))
            )
            candidates_created += 1
        except Exception as e:
            print(f"创建候选人失败: {e}")
            continue
    
    print(f"创建了 {candidates_created} 个候选人")
    
    # 创建申请记录
    all_positions = list(JobPosition.objects.filter(status='published'))
    all_candidates = list(Candidate.objects.all())
    
    applications_created = 0
    used_combinations = set()
    for i in range(40):
        try:
            candidate = random.choice(all_candidates)
            position = random.choice(all_positions)
            combination = (candidate.id, position.id)
            
            # 避免重复的候选人-职位组合
            if combination in used_combinations:
                continue
            used_combinations.add(combination)
            
            application = JobApplication.objects.create(
                candidate=candidate,
                position=position,
                cover_letter=f"申请信内容{i+1}",
                status=random.choice(['submitted', 'reviewing', 'interview', 'offer', 'hired', 'rejected']),
                rating=random.randint(1, 5),
                reviewer=User.objects.filter(is_staff=True).first()
            )
            applications_created += 1
        except Exception as e:
            print(f"创建申请失败: {e}")
            continue
    
    print(f"创建了 {applications_created} 个申请")
    
    # 5. 创建考勤数据
    # 创建请假类型
    leave_types_data = [
        {'name': '年假', 'code': 'ANNUAL', 'max_days_per_year': 10},
        {'name': '病假', 'code': 'SICK', 'max_days_per_year': 30},
        {'name': '事假', 'code': 'PERSONAL', 'max_days_per_year': 5},
        {'name': '婚假', 'code': 'MARRIAGE', 'max_days_per_year': 3},
        {'name': '产假', 'code': 'MATERNITY', 'max_days_per_year': 98},
    ]
    
    for leave_data in leave_types_data:
        LeaveType.objects.get_or_create(
            code=leave_data['code'],
            defaults={
                'name': leave_data['name'],
                'max_days_per_year': leave_data['max_days_per_year'],
                'need_approval': True,
                'is_active': True
            }
        )
    
    # 创建考勤记录
    attendance_created = 0
    for employee in all_employees[:20]:  # 为前20个员工创建考勤记录
        for days_back in range(30):  # 最近30天
            record_date = date.today() - timedelta(days=days_back)
            if record_date.weekday() < 5:  # 工作日
                try:
                    # 签到记录
                    AttendanceRecord.objects.create(
                        employee=employee,
                        date=record_date,
                        record_type='check_in',
                        time=datetime.combine(record_date, datetime.min.time().replace(
                            hour=random.randint(8, 9),
                            minute=random.randint(0, 59)
                        )).time(),
                        location='公司'
                    )
                    
                    # 签退记录
                    AttendanceRecord.objects.create(
                        employee=employee,
                        date=record_date,
                        record_type='check_out',
                        time=datetime.combine(record_date, datetime.min.time().replace(
                            hour=random.randint(17, 19),
                            minute=random.randint(0, 59)
                        )).time(),
                        location='公司'
                    )
                    attendance_created += 2
                    
                except Exception as e:
                    print(f"创建考勤记录失败: {e}")
                    continue
    
    print(f"创建了 {attendance_created} 条考勤记录")
    
    print("测试数据生成完成！")
    print("\n=== 数据统计 ===")
    print(f"员工总数: {Employee.objects.count()}")
    print(f"部门总数: {Department.objects.count()}")
    print(f"薪资记录总数: {SalaryRecord.objects.count()}")
    print(f"职位总数: {JobPosition.objects.count()}")
    print(f"候选人总数: {Candidate.objects.count()}")
    print(f"申请总数: {JobApplication.objects.count()}")
    print(f"考勤记录总数: {AttendanceRecord.objects.count()}")

if __name__ == '__main__':
    create_test_data()