X-Financial/server/tests/test_hierarchical_expense_memory_foundation.py

from __future__ import annotations

from datetime import UTC, datetime, timedelta
from decimal import Decimal

import pytest
from sqlalchemy import create_engine
from sqlalchemy.exc import IntegrityError
from sqlalchemy.orm import Session, sessionmaker
from sqlalchemy.pool import StaticPool

from app.api.deps import _authenticate_bearer_user
from app.db.base import Base
from app.models.ai_memory import MemoryEntry
from app.models.employee import Employee
from app.services.auth import AuthService
from app.services.auth_sessions import AuthSessionService
from app.services.employee import EmployeeService


def _build_session() -> 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 _memory_entry(
    *,
    scope_type: str,
    scope_id: str,
    tenant_id: str = "tenant-a",
    origin_type: str = "learned",
    with_management_audit: bool = False,
) -> MemoryEntry:
    now = datetime.now(UTC)
    return MemoryEntry(
        tenant_id=tenant_id,
        scope_type=scope_type,
        scope_id=scope_id,
        origin_type=origin_type,
        managed_by="admin@example.com" if with_management_audit else None,
        managed_at=now if with_management_audit else None,
        management_reason="统一差旅交通基线" if with_management_audit else None,
        policy_version="travel-policy-v1" if with_management_audit else None,
        value_json={"value": "火车"},
        value_fingerprint=f"{scope_type}:{scope_id}",
        status="active",
        evidence_count=0,
        approved_evidence_count=0,
        confidence=Decimal("1.0000"),
        last_evidence_at=now,
        candidate_expires_at=now + timedelta(days=90),
        activated_at=now,
        active_expires_at=now + timedelta(days=180),
        created_at=now,
        updated_at=now,
    )


def test_memory_entry_accepts_learned_user_and_audited_organization_scopes() -> None:
    with _build_session() as db:
        learned_entry = _memory_entry(
            scope_type="user",
            scope_id="employee-a",
        )
        learned_entry.policy_version = "expense-application-memory.v1"
        db.add(learned_entry)
        db.add(
            _memory_entry(
                scope_type="department",
                scope_id="department-a",
                origin_type="admin_managed",
                with_management_audit=True,
            )
        )
        db.add(
            _memory_entry(
                scope_type="enterprise",
                scope_id="tenant-a",
                origin_type="admin_managed",
                with_management_audit=True,
            )
        )

        db.commit()

        assert db.query(MemoryEntry).count() == 3


@pytest.mark.parametrize(
    ("entry", "expected_constraint"),
    [
        (
            _memory_entry(
                scope_type="department",
                scope_id="department-a",
                origin_type="learned",
            ),
            "ck_memory_entries_scope_origin",
        ),
        (
            _memory_entry(
                scope_type="enterprise",
                scope_id="another-tenant",
                origin_type="admin_managed",
                with_management_audit=True,
            ),
            "ck_memory_entries_enterprise_scope",
        ),
        (
            _memory_entry(
                scope_type="department",
                scope_id="department-a",
                origin_type="admin_managed",
            ),
            "ck_memory_entries_management_audit",
        ),
    ],
)
def test_memory_entry_rejects_invalid_scope_origin_and_management_audit(
    entry: MemoryEntry,
    expected_constraint: str,
) -> None:
    with _build_session() as db:
        db.add(entry)

        with pytest.raises(IntegrityError, match=expected_constraint):
            db.commit()


def test_admin_managed_memory_rejects_partially_missing_audit_fields() -> None:
    with _build_session() as db:
        entry = _memory_entry(
            scope_type="department",
            scope_id="department-a",
            origin_type="admin_managed",
            with_management_audit=True,
        )
        entry.management_reason = "   "
        db.add(entry)

        with pytest.raises(IntegrityError, match="ck_memory_entries_management_audit"):
            db.commit()


def test_authenticated_session_restores_tenant_and_stable_department_id() -> None:
    with _build_session() as db:
        employee_snapshot = EmployeeService(db).list_employees()[0]
        employee = db.get(Employee, employee_snapshot.id)
        assert employee is not None
        authenticated_at_login = AuthService(db)._build_employee_user(employee)
        authenticated_at_login.tenant_id = "tenant-session-a"
        access_token, auth_session = AuthSessionService(db).issue(
            authenticated_at_login,
            metric_session_id="metric-session-a",
        )
        db.commit()

        restored_user = AuthService(db).get_session_user(auth_session)
        current_user = _authenticate_bearer_user(db, f"Bearer {access_token}")

        assert restored_user is not None
        assert restored_user.tenant_id == "tenant-session-a"
        assert restored_user.department_id == employee.organization_unit_id
        assert current_user.tenant_id == "tenant-session-a"
        assert current_user.department_id == employee.organization_unit_id
        assert current_user.department_name == employee.organization_unit.name


def test_platform_admin_has_no_department_scope() -> None:
    with _build_session() as db:
        record = type(
            "AdminRecord",
            (),
            {"account": "admin", "email": "admin@example.com"},
        )()
        admin_user = AuthService(db)._build_admin_user(record)

        assert admin_user.department_id is None