company-ai-platform/tests/test_v2_hardening.py

import json
from pathlib import Path

import pytest
from sqlalchemy import create_engine, select
from sqlalchemy.orm import Session, sessionmaker

from app.core.background.task_queue.reports import _queue_celery_report_push
from app.core.config import Settings
from app.core.http.masking import MASKED_VALUE, mask_sensitive
from app.core.utils.time import utc_now
from app.modules.ai_memory.constants import AIMemoryStatus
from app.modules.ai_memory.models import AIMemoryEntry
from app.modules.ai_memory.service import AIMemoryService
from app.modules.audit.constants import AUDIT_REDACTED_VALUE
from app.modules.audit.models import AuditLog
from app.modules.audit.schemas import AuditLogCreate
from app.modules.audit.service import AuditService
from app.modules.business.models import Project, ReportPushRun
from app.modules.events.models import DomainEvent
from app.modules.feishu.service import FeishuService
from app.modules.reports.constants import ReportPushStatus
from app.tasks import celery_app
from app.tasks.reports import _push_report
from app.tools import init_db


def _production_settings(**overrides: object) -> Settings:
    values: dict[str, object] = {
        "app_env": "production",
        "database_url": (
            "postgresql+psycopg://app:runtime-database-password-2026@db/app"
        ),
        "api_key": "runtime-service-key-2026-primary",
        "audit_api_key": "runtime-audit-key-2026-independent",
        "cors_origins": ["https://internal.example.com"],
        "debug": False,
        "mask_sensitive_responses": True,
        "read_only_mode": True,
        "feishu_event_transport": "disabled",
    }
    values.update(overrides)
    return Settings(_env_file=None, **values)


def test_production_keys_must_be_enabled_isolated_and_safe() -> None:
    settings = _production_settings()
    assert settings.api_key == "runtime-service-key-2026-primary"

    with pytest.raises(ValueError, match="API_KEY or API_KEYS"):
        _production_settings(
            api_key=None,
            api_keys=[{"key": "disabled-key", "enabled": False}],
        )

    with pytest.raises(ValueError, match="cannot overlap"):
        _production_settings(audit_api_key="runtime-service-key-2026-primary")

    with pytest.raises(ValueError, match="MASK_SENSITIVE_RESPONSES"):
        _production_settings(mask_sensitive_responses=False)

    with pytest.raises(ValueError, match="READ_ONLY_MODE"):
        _production_settings(read_only_mode=False)


def test_audit_json_strings_and_responses_mask_sensitive_fields() -> None:
    engine = create_engine("sqlite://")
    AuditLog.__table__.create(engine)
    try:
        with Session(engine) as db:
            service = AuditService(db)
            dict_record = service.log(
                AuditLogCreate(
                    action="security.redaction.dict",
                    request_payload={
                        "accessToken": "access-secret",
                        "nested": {
                            "clientSecret": "client-secret",
                            "safe": "visible",
                        },
                    },
                )
            )
            json_record = service.log(
                AuditLogCreate(
                    action="security.redaction.json",
                    request_payload=json.dumps(
                        {
                            "authorization": "Bearer secret",
                            "nested": {
                                "refreshToken": "refresh-secret",
                                "safe": "visible",
                            },
                        }
                    ),
                )
            )
            dict_stored = json.loads(dict_record.request_payload or "{}")
            json_stored = json.loads(json_record.request_payload or "{}")

        assert dict_stored["accessToken"] == AUDIT_REDACTED_VALUE
        assert dict_stored["nested"]["clientSecret"] == AUDIT_REDACTED_VALUE
        assert dict_stored["nested"]["safe"] == "visible"
        assert json_stored["authorization"] == AUDIT_REDACTED_VALUE
        assert json_stored["nested"]["refreshToken"] == AUDIT_REDACTED_VALUE
        assert json_stored["nested"]["safe"] == "visible"

        masked = mask_sensitive(
            {
                "accessToken": "access-secret",
                "nested": {
                    "clientSecret": "client-secret",
                    "privateKey": "private-secret",
                    "safe": "visible",
                },
            }
        )
        assert masked["accessToken"] == MASKED_VALUE
        assert masked["nested"]["clientSecret"] == MASKED_VALUE
        assert masked["nested"]["privateKey"] == MASKED_VALUE
        assert masked["nested"]["safe"] == "visible"
    finally:
        engine.dispose()


def test_report_queue_failure_is_persisted(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    engine = create_engine("sqlite://")
    ReportPushRun.__table__.create(engine)
    factory = sessionmaker(bind=engine, expire_on_commit=False)

    class BrokenSignature:
        def apply_async(self, task_id: str) -> None:
            _ = task_id
            raise RuntimeError("broker unavailable")

    monkeypatch.setattr("app.core.database.SessionLocal", factory)
    monkeypatch.setattr(
        celery_app,
        "signature",
        lambda *args, **kwargs: BrokenSignature(),
    )
    try:
        with pytest.raises(RuntimeError, match="broker unavailable"):
            _queue_celery_report_push(
                task_name="reports.test",
                report_type="test",
                title="Test report",
                receive_id=None,
                receive_id_type="chat_id",
                actor="pytest",
            )

        with factory() as db:
            run = db.execute(select(ReportPushRun)).scalar_one()
            assert run.status == ReportPushStatus.FAILED
            assert run.task_id
            assert run.error_message == "broker unavailable"
    finally:
        engine.dispose()


def test_eager_report_terminal_state_is_not_overwritten(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    engine = create_engine("sqlite://")
    ReportPushRun.__table__.create(engine)
    factory = sessionmaker(bind=engine, expire_on_commit=False)

    class EagerSignature:
        def __init__(self, push_run_code: str):
            self.push_run_code = push_run_code

        def apply_async(self, task_id: str) -> None:
            with factory() as db:
                run = db.execute(
                    select(ReportPushRun).where(
                        ReportPushRun.code == self.push_run_code
                    )
                ).scalar_one()
                assert run.task_id == task_id
                run.status = ReportPushStatus.SUCCESS
                db.commit()

    def eager_signature(*args: object, **kwargs: object) -> EagerSignature:
        _ = args
        task_kwargs = kwargs["kwargs"]
        assert isinstance(task_kwargs, dict)
        return EagerSignature(str(task_kwargs["push_run_code"]))

    monkeypatch.setattr("app.core.database.SessionLocal", factory)
    monkeypatch.setattr(celery_app, "signature", eager_signature)
    try:
        queued = _queue_celery_report_push(
            task_name="reports.test",
            report_type="test",
            title="Test report",
            receive_id=None,
            receive_id_type="chat_id",
            actor="pytest",
        )

        with factory() as db:
            run = db.execute(select(ReportPushRun)).scalar_one()
            assert run.status == ReportPushStatus.SUCCESS
            assert run.task_id == queued["task_id"]
    finally:
        engine.dispose()


def test_successful_report_task_redelivery_does_not_send_twice(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    engine = create_engine("sqlite://")
    for table in (AuditLog.__table__, DomainEvent.__table__, ReportPushRun.__table__):
        table.create(engine)
    factory = sessionmaker(bind=engine, expire_on_commit=False)
    calls = {"build": 0, "send": 0}

    def build_report(service: object, actor: str) -> dict[str, object]:
        _ = service, actor
        calls["build"] += 1
        return {
            "title": "Idempotent report",
            "report_type": "test",
            "lines": ["ok"],
            "content": "ok",
        }

    def fake_send_card(
        service: object,
        card: dict,
        receive_id: str | None,
        receive_id_type: str,
        actor: str,
    ) -> dict[str, object]:
        _ = service, card, receive_id, receive_id_type, actor
        calls["send"] += 1
        return {"code": 0, "message_id": "om-idempotent"}

    monkeypatch.setattr("app.tasks.reports.SessionLocal", factory)
    monkeypatch.setattr(FeishuService, "send_card", fake_send_card)
    try:
        with factory() as db:
            run = ReportPushRun(
                code="PUSH-IDEMPOTENT",
                report_type="test",
                title="Idempotent report",
                receive_id="oc-idempotent",
                receive_id_type="chat_id",
                status=ReportPushStatus.QUEUED,
                actor="pytest",
            )
            db.add(run)
            db.commit()

        first = _push_report(build_report, "oc-idempotent", "chat_id", "pytest", run.code)
        second = _push_report(build_report, "oc-idempotent", "chat_id", "pytest", run.code)

        assert first == second == {"code": 0, "message_id": "om-idempotent"}
        assert calls == {"build": 1, "send": 1}
        with factory() as db:
            stored = db.execute(select(ReportPushRun)).scalar_one()
            assert stored.status == ReportPushStatus.SUCCESS
            assert stored.sent_at is not None
    finally:
        engine.dispose()


def test_init_db_upgrades_to_alembic_head(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    called: dict[str, object] = {}

    def fake_upgrade(config: object, revision: str) -> None:
        called["config"] = config
        called["revision"] = revision

    monkeypatch.setattr(init_db.command, "upgrade", fake_upgrade)
    init_db.main()

    config = called["config"]
    assert called["revision"] == "head"
    assert Path(config.config_file_name).name == "alembic.ini"
    assert Path(config.get_main_option("script_location")).name == "alembic"


def test_memory_retention_does_not_commit_caller_transaction() -> None:
    engine = create_engine("sqlite://")
    AIMemoryEntry.__table__.create(engine)
    Project.__table__.create(engine)
    factory = sessionmaker(bind=engine, expire_on_commit=False)
    try:
        with factory() as db:
            expired = AIMemoryEntry(
                code="MEM-EXPIRED",
                scope="project",
                subject="P-ROLLBACK",
                content="expired",
                status=AIMemoryStatus.ACTIVE,
                expires_at=utc_now(),
            )
            db.add(expired)
            db.commit()

            db.add(Project(code="P-ROLLBACK", name="Must roll back"))
            assert AIMemoryService(db).list_entries(scope="project") == []
            db.rollback()

        with factory() as db:
            assert db.execute(
                select(Project).where(Project.code == "P-ROLLBACK")
            ).scalar_one_or_none() is None
            stored = db.execute(
                select(AIMemoryEntry).where(AIMemoryEntry.code == "MEM-EXPIRED")
            ).scalar_one()
            assert stored.status == AIMemoryStatus.ACTIVE
    finally:
        engine.dispose()