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"] == "销售部"