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.models.tenant import Tenant 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: db.add( Tenant( tenant_id="tenant-session-a", tenant_code="tenant-session-a", name="会话租户 A", status="active", ) ) db.flush() employee_snapshot = EmployeeService( db, tenant_id="tenant-session-a", ).list_employees()[0] employee = db.get(Employee, employee_snapshot.id) assert employee is not None authenticated_at_login = AuthService(db)._build_employee_user(employee) 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