company-ai-platform/tests/test_feishu_long_connection_health.py

from types import SimpleNamespace
from typing import Any

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

from app.core.config import get_settings
from app.core.database import Base, get_db
from app.core.database.migrations import expected_alembic_heads
from app.modules.feishu import long_connection
from app.modules.observability.constants import (
    HeartbeatComponent,
    HeartbeatStatus,
    ObservabilityKey,
    ObservabilityStatus,
)
from app.modules.observability.routes import router as observability_router
from app.modules.observability.service import ObservabilityService


def test_connection_heartbeat_tracks_disconnect_and_recovery(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    client = SimpleNamespace(_conn=None)
    statuses: list[str] = []

    class FakeStopEvent:
        def __init__(self) -> None:
            self.wait_count = 0

        def is_set(self) -> bool:
            return self.wait_count >= 3

        def wait(self, timeout: float) -> None:
            assert timeout == 1
            self.wait_count += 1
            if self.wait_count == 1:
                client._conn = SimpleNamespace(
                    state=SimpleNamespace(name="OPEN"),
                )
            elif self.wait_count == 2:
                client._conn = SimpleNamespace(
                    state=SimpleNamespace(name="CLOSED"),
                )

    monkeypatch.setattr(
        long_connection,
        "_record_runtime_heartbeat",
        lambda _instance_id, status_value: statuses.append(status_value),
    )

    long_connection._heartbeat_loop(
        FakeStopEvent(),  # type: ignore[arg-type]
        "feishu-events-test",
        30,
        client,
    )

    assert statuses == [
        HeartbeatStatus.DEGRADED,
        HeartbeatStatus.OK,
        HeartbeatStatus.DEGRADED,
    ]


def test_connection_adapter_supports_sdk_state_variants() -> None:
    assert long_connection._sdk_connection_is_open(SimpleNamespace(_conn=None)) is False
    unknown_client = SimpleNamespace(_conn=SimpleNamespace())
    assert long_connection._sdk_connection_is_open(unknown_client) is False
    assert (
        long_connection._connection_heartbeat_status(unknown_client)
        == HeartbeatStatus.DEGRADED
    )
    assert (
        long_connection._sdk_connection_is_open(
            SimpleNamespace(_conn=SimpleNamespace(closed=False))
        )
        is True
    )
    assert (
        long_connection._sdk_connection_is_open(
            SimpleNamespace(_conn=SimpleNamespace(open=False))
        )
        is False
    )
    assert (
        long_connection._sdk_connection_is_open(
            SimpleNamespace(
                _conn=SimpleNamespace(state=SimpleNamespace(name="OPEN"))
            )
        )
        is True
    )
    assert (
        long_connection._sdk_connection_is_open(
            SimpleNamespace(
                _conn=SimpleNamespace(state=SimpleNamespace(name="CLOSING"))
            )
        )
        is False
    )


def test_readiness_returns_503_while_disconnected_and_recovers(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    monkeypatch.setenv("APP_ENV", "test")
    monkeypatch.setenv("FEISHU_EVENT_TRANSPORT", "long_connection")
    monkeypatch.setenv("FEISHU_APP_ID", "app-id")
    monkeypatch.setenv("FEISHU_APP_SECRET", "app-secret")
    monkeypatch.setenv("FEISHU_USER_FEATURES_ENABLED", "false")
    monkeypatch.setenv("SCHEDULER_ENABLED", "false")
    monkeypatch.setenv("TASK_QUEUE_ENABLED", "false")
    get_settings.cache_clear()

    engine = create_engine(
        "sqlite://",
        connect_args={"check_same_thread": False},
        poolclass=StaticPool,
    )
    Base.metadata.create_all(engine)
    with engine.begin() as connection:
        connection.execute(
            text("create table alembic_version (version_num varchar(32))")
        )
        connection.execute(
            text("insert into alembic_version (version_num) values (:revision)"),
            {"revision": expected_alembic_heads()[0]},
        )
    session_factory = sessionmaker(bind=engine, expire_on_commit=False)
    app = FastAPI()
    app.include_router(observability_router, prefix="/api/v1")

    def override_get_db() -> Any:
        db = session_factory()
        try:
            yield db
        finally:
            db.close()

    app.dependency_overrides[get_db] = override_get_db
    client = TestClient(app)
    try:
        with Session(engine) as db:
            ObservabilityService(db).record_heartbeat(
                HeartbeatComponent.FEISHU_EVENTS,
                "feishu-events-test",
                status_value=HeartbeatStatus.DEGRADED,
            )

        disconnected = client.get("/api/v1/health/ready")

        assert disconnected.status_code == 503
        disconnected_check = disconnected.json()[ObservabilityKey.CHECKS][
            ObservabilityKey.FEISHU_EVENTS
        ]
        assert (
            disconnected_check[ObservabilityKey.STATUS]
            == ObservabilityStatus.DEGRADED
        )
        assert disconnected_check["reason"] == "heartbeat_degraded"

        with Session(engine) as db:
            ObservabilityService(db).record_heartbeat(
                HeartbeatComponent.FEISHU_EVENTS,
                "feishu-events-test",
                status_value=HeartbeatStatus.OK,
            )

        recovered = client.get("/api/v1/health/ready")

        assert recovered.status_code == 200
        recovered_check = recovered.json()[ObservabilityKey.CHECKS][
            ObservabilityKey.FEISHU_EVENTS
        ]
        assert recovered_check[ObservabilityKey.STATUS] == ObservabilityStatus.OK
        assert recovered_check["reason"] is None
    finally:
        client.close()
        engine.dispose()
        get_settings.cache_clear()


@pytest.mark.parametrize(
    ("configured_identities", "expected_reason"),
    [
        ([""], "admin_identities_missing"),
        (["not-an-identity"], "admin_identities_invalid"),
    ],
)
def test_production_readiness_parses_admin_identities_without_exposing_them(
    monkeypatch: pytest.MonkeyPatch,
    configured_identities: list[str],
    expected_reason: str,
) -> None:
    sensitive_value = configured_identities[0]
    settings = SimpleNamespace(
        app_env="production",
        feishu_user_features_enabled=True,
        feishu_admin_identities=configured_identities,
        feishu_event_transport="long_connection",
        feishu_app_id="app-id",
        feishu_app_secret="app-secret",
        feishu_app_type="self",
        feishu_default_tenant_key=None,
    )
    monkeypatch.setattr(
        "app.modules.observability.service.get_settings",
        lambda: settings,
    )
    engine = create_engine("sqlite://")
    Base.metadata.create_all(engine)
    try:
        with Session(engine) as db:
            result = ObservabilityService(db)._feishu_subscriptions_check()

        assert result[ObservabilityKey.STATUS] == ObservabilityStatus.DEGRADED
        assert expected_reason in result["reasons"]
        if sensitive_value:
            assert sensitive_value not in str(result)
    finally:
        engine.dispose()