from __future__ import annotations from datetime import UTC, date, datetime from sqlalchemy import create_engine, func, select from sqlalchemy.orm import Session, sessionmaker from sqlalchemy.pool import StaticPool from app.db.base import Base from app.models.budget import BudgetAllocation, BudgetReservation, BudgetTransaction from app.models.employee import Employee from app.models.financial_record import ExpenseClaim, ExpenseClaimItem from app.models.organization import OrganizationUnit from app.models.risk_observation import RiskObservation from app.services.budget import BudgetService from app.services.demo_company_simulation_seed import ( SIM_CLAIM_PREFIX, SIM_EMPLOYEE_PREFIX, HalfYearExpenseSimulationSeeder, SimulationConfig, ) def build_session() -> Session: engine = create_engine( "sqlite+pysqlite:///:memory:", connect_args={"check_same_thread": False}, poolclass=StaticPool, ) Base.metadata.create_all(bind=engine) session_factory = sessionmaker(bind=engine, autoflush=False, autocommit=False) return session_factory() def seed_company(db: Session) -> None: tech = OrganizationUnit( id="dept-tech", unit_code="TECH-DEPT", name="技术部", unit_type="department", cost_center="CC-6100", location="北京", ) market = OrganizationUnit( id="dept-market", unit_code="MARKET-DEPT", name="市场部", unit_type="department", cost_center="CC-4100", location="上海", ) db.add_all([tech, market]) for index in range(3): db.add( Employee( id=f"emp-existing-{index}", employee_no=f"E-EXISTING-{index}", name=f"现有员工{index}", email=f"existing-{index}@xf.com", grade="P5", position="主管", organization_unit=tech if index % 2 == 0 else market, cost_center="CC-6100" if index % 2 == 0 else "CC-4100", ) ) db.commit() def test_half_year_simulation_preview_and_apply_are_idempotent() -> None: with build_session() as db: seed_company(db) config = SimulationConfig(target_employees=8, start_date=date(2026, 1, 1), months=6, seed=7) preview = HalfYearExpenseSimulationSeeder(db, config).preview() assert preview.mode == "dry-run" assert preview.current_employee_count == 3 assert preview.employees_to_create == 5 assert preview.claims_to_create >= 24 assert preview.budget_allocations_to_create > 0 assert preview.budget_transactions_to_create > 0 applied = HalfYearExpenseSimulationSeeder(db, config).apply() db.commit() assert applied.mode == "apply" assert applied.employees_to_create == 5 assert db.scalar(select(func.count()).select_from(Employee)) == 8 assert db.scalar(select(func.count()).select_from(ExpenseClaim)) == applied.claims_to_create assert ( db.scalar(select(func.count()).select_from(ExpenseClaimItem)) == applied.claim_items_to_create ) assert ( db.scalar(select(func.count()).select_from(BudgetAllocation)) == applied.budget_allocations_to_create ) assert ( db.scalar(select(func.count()).select_from(BudgetTransaction)) == applied.budget_transactions_to_create ) assert ( db.scalar(select(func.count()).select_from(BudgetReservation)) == applied.budget_reservations_to_create ) assert ( db.scalar(select(func.count()).select_from(RiskObservation)) == applied.risk_observations_to_create ) repeated = HalfYearExpenseSimulationSeeder(db, config).apply() db.commit() assert repeated.employees_to_create == 0 assert repeated.claims_to_create == 0 assert repeated.budget_allocations_to_create == 0 assert repeated.budget_transactions_to_create == 0 assert repeated.budget_reservations_to_create == 0 assert repeated.risk_observations_to_create == 0 def test_half_year_simulation_feeds_budget_summary() -> None: with build_session() as db: seed_company(db) config = SimulationConfig( target_employees=10, start_date=date(2026, 1, 1), months=6, seed=11, ) HalfYearExpenseSimulationSeeder(db, config).apply() db.commit() summary = BudgetService(db).get_summary(fiscal_year=2026, period_key="2026Q2") sim_claim_count = db.scalar( select(func.count()).select_from(ExpenseClaim).where(ExpenseClaim.claim_no.like(f"{SIM_CLAIM_PREFIX}%")) ) sim_employee_count = db.scalar( select(func.count()).select_from(Employee).where(Employee.employee_no.like(f"{SIM_EMPLOYEE_PREFIX}%")) ) assert sim_claim_count and sim_claim_count >= 30 assert sim_employee_count == 7 assert summary.trend assert {item.period_key for item in summary.trend} == {"2026Q1", "2026Q2"} assert summary.warning_count + summary.over_budget_count > 0 def test_half_year_simulation_excludes_admin_and_visible_month_has_real_volume() -> None: with build_session() as db: seed_company(db) db.add( Employee( id="emp-admin", employee_no="ADMIN", name="admin", email="admin@xf.com", grade="P8", position="admin", ) ) db.commit() config = SimulationConfig( target_employees=100, start_date=date(2026, 1, 1), months=6, seed=20260602, ) HalfYearExpenseSimulationSeeder(db, config).apply() db.commit() admin_claim_count = db.scalar( select(func.count()) .select_from(ExpenseClaim) .where(ExpenseClaim.employee_name == "admin") ) visible_claim_count = db.scalar( select(func.count()) .select_from(ExpenseClaim) .where(ExpenseClaim.claim_no.like(f"{SIM_CLAIM_PREFIX}%")) .where(ExpenseClaim.occurred_at >= datetime(2026, 6, 1, tzinfo=UTC)) .where(ExpenseClaim.occurred_at < datetime(2026, 6, 3, tzinfo=UTC)) ) earliest_claim_day = db.scalar( select(func.min(ExpenseClaim.occurred_at)).where( ExpenseClaim.claim_no.like(f"{SIM_CLAIM_PREFIX}%") ) ) latest_claim_day = db.scalar( select(func.max(ExpenseClaim.occurred_at)).where( ExpenseClaim.claim_no.like(f"{SIM_CLAIM_PREFIX}%") ) ) assert admin_claim_count == 0 assert visible_claim_count is not None assert 400 <= visible_claim_count <= 500 assert earliest_claim_day is not None assert latest_claim_day is not None assert earliest_claim_day.date() >= date(2026, 1, 1) assert latest_claim_day.date() <= date(2026, 6, 2)