X-Financial/server/tests/test_budget_endpoints.py

from __future__ import annotations

from collections.abc import Generator
from datetime import UTC, datetime
from decimal import Decimal

from fastapi.testclient import TestClient
from sqlalchemy import create_engine
from sqlalchemy.orm import Session, sessionmaker
from sqlalchemy.pool import StaticPool

from app.api.deps import get_db
from app.db.base import Base
from app.main import create_app
from app.models.budget import BudgetAllocation


def build_session_factory() -> sessionmaker[Session]:
    engine = create_engine(
        "sqlite+pysqlite:///:memory:",
        connect_args={"check_same_thread": False},
        poolclass=StaticPool,
    )
    Base.metadata.create_all(bind=engine)
    return sessionmaker(bind=engine, autoflush=False, autocommit=False)


def build_client() -> tuple[TestClient, sessionmaker[Session]]:
    session_factory = build_session_factory()
    app = create_app()

    def override_db() -> Generator[Session, None, None]:
        db = session_factory()
        try:
            yield db
        finally:
            db.close()

    app.dependency_overrides[get_db] = override_db
    return TestClient(app), session_factory


def seed_budget_allocations(db: Session) -> None:
    now = datetime.now(UTC)
    db.add_all(
        [
            BudgetAllocation(
                id="budget-market-travel",
                budget_no="BUD-MARKET-TRAVEL",
                fiscal_year=2026,
                period_type="quarter",
                period_key="2026Q2",
                department_id="dept-market",
                department_name="市场部",
                cost_center="CC-4100",
                project_code=None,
                subject_code="travel",
                subject_name="差旅费",
                original_amount=Decimal("50000.00"),
                adjusted_amount=Decimal("0.00"),
                status="active",
                warning_threshold=Decimal("80.00"),
                control_action="block",
                created_at=now,
                updated_at=now,
            ),
            BudgetAllocation(
                id="budget-finance-office",
                budget_no="BUD-FINANCE-OFFICE",
                fiscal_year=2026,
                period_type="quarter",
                period_key="2026Q2",
                department_id="dept-finance",
                department_name="财务部",
                cost_center="CC-2100",
                project_code=None,
                subject_code="office",
                subject_name="办公费",
                original_amount=Decimal("30000.00"),
                adjusted_amount=Decimal("0.00"),
                status="active",
                warning_threshold=Decimal("80.00"),
                control_action="block",
                created_at=now,
                updated_at=now,
            ),
            BudgetAllocation(
                id="budget-market-software-hidden",
                budget_no="BUD-MARKET-SOFTWARE",
                fiscal_year=2026,
                period_type="quarter",
                period_key="2026Q2",
                department_id="dept-market",
                department_name="市场部",
                cost_center="CC-4100",
                project_code=None,
                subject_code="software",
                subject_name="软件服务费",
                original_amount=Decimal("90000.00"),
                adjusted_amount=Decimal("0.00"),
                status="active",
                warning_threshold=Decimal("80.00"),
                control_action="block",
                created_at=now,
                updated_at=now,
            ),
        ]
    )
    db.commit()


def test_admin_can_view_all_budget_allocations_without_is_admin_header() -> None:
    client, session_factory = build_client()
    with session_factory() as db:
        seed_budget_allocations(db)

    response = client.get(
        "/api/v1/budgets/summary",
        headers={"x-auth-username": "admin", "x-auth-role-codes": "manager"},
    )

    assert response.status_code == 200
    payload = response.json()
    assert {item["department_name"] for item in payload["allocations"]} == {"市场部", "财务部"}
    assert {item["subject_code"] for item in payload["allocations"]} == {"travel", "office"}


def test_budget_monitor_is_limited_to_own_department_scope() -> None:
    client, session_factory = build_client()
    with session_factory() as db:
        seed_budget_allocations(db)

    response = client.get(
        "/api/v1/budgets/summary?cost_center=CC-2100",
        headers={
            "x-auth-username": "monitor@example.com",
            "x-auth-role-codes": "budget_monitor",
            "x-auth-cost-center": "CC-4100",
        },
    )

    assert response.status_code == 200
    payload = response.json()
    assert [item["cost_center"] for item in payload["allocations"]] == ["CC-4100"]
    assert [item["subject_code"] for item in payload["allocations"]] == ["travel"]


def test_finance_user_cannot_enter_budget_center() -> None:
    client, session_factory = build_client()
    with session_factory() as db:
        seed_budget_allocations(db)

    response = client.get(
        "/api/v1/budgets/summary",
        headers={"x-auth-username": "finance@example.com", "x-auth-role-codes": "finance"},
    )

    assert response.status_code == 403


def test_budget_monitor_cannot_edit_and_admin_can_edit() -> None:
    client, session_factory = build_client()
    with session_factory() as db:
        seed_budget_allocations(db)

    payload = {
        "fiscal_year": 2026,
        "period_type": "quarter",
        "period_key": "2026Q2",
        "department_id": "dept-sales",
        "department_name": "销售部",
        "cost_center": "CC-5100",
        "project_code": None,
        "subject_code": "travel",
        "subject_name": "差旅费",
        "original_amount": "20000.00",
        "warning_threshold": "80.00",
        "control_action": "block",
        "description": "销售部差旅预算",
    }

    monitor_response = client.post(
        "/api/v1/budgets/allocations",
        json=payload,
        headers={
            "x-auth-username": "monitor@example.com",
            "x-auth-role-codes": "budget_monitor",
            "x-auth-cost-center": "CC-4100",
        },
    )
    assert monitor_response.status_code == 403

    admin_response = client.post(
        "/api/v1/budgets/allocations",
        json=payload,
        headers={"x-auth-username": "admin", "x-auth-role-codes": "manager"},
    )
    assert admin_response.status_code == 201
    assert admin_response.json()["department_name"] == "销售部"