roboco/tests/unit/api/test_a2a_subscribe.py

"""SSE ``subscribe_to_task`` is authenticated like the rest of the a2a
message surface and opens a SHORT-LIVED DB session per poll iteration
(via ``get_session_factory()``) instead of holding the request-scoped
``db: DbSession`` for the full SSE lifetime, which exhausted the asyncpg
pool one connection per connected client.
"""

from __future__ import annotations

from typing import TYPE_CHECKING, Any, cast
from unittest.mock import AsyncMock, MagicMock

import pytest
from fastapi import FastAPI
from httpx import ASGITransport, AsyncClient
from roboco.agents_config import issue_agent_token
from roboco.api.routes import a2a as a2a_module
from roboco.api.routes.a2a import router as a2a_router
from roboco.db.base import get_db

if TYPE_CHECKING:
    from collections.abc import AsyncIterator

    from fastapi.routing import APIRoute

_SECRET = "test-secret-for-a2a-subscribe"
_AGENT_ID = "00000000-0000-0000-0000-000000000003"
_HTTP_200 = 200
_HTTP_401 = 401
_HTTP_404 = 404


@pytest.fixture
async def a2a_client() -> AsyncIterator[AsyncClient]:
    app = FastAPI()
    app.include_router(a2a_router, prefix="/api/a2a")
    async with AsyncClient(
        transport=ASGITransport(app=app), base_url="http://test"
    ) as client:
        yield client
    app.dependency_overrides.clear()


# ---------------------------------------------------------------------------
# Auth gate (F023 parity)
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_subscribe_rejects_missing_token_when_required(
    a2a_client: AsyncClient, monkeypatch: pytest.MonkeyPatch
) -> None:
    """Strict mode + no X-Agent-Token => 401, never reaches the generator."""
    monkeypatch.setenv("ROBOCO_AGENT_AUTH_SECRET", _SECRET)
    monkeypatch.setenv("ROBOCO_AGENT_AUTH_REQUIRED", "true")
    r = await a2a_client.get(
        "/api/a2a/tasks/some-task/subscribe",
        headers={"X-Agent-ID": _AGENT_ID, "X-Agent-Role": "developer"},
    )
    assert r.status_code == _HTTP_401


@pytest.mark.asyncio
async def test_subscribe_rejects_forged_token_even_in_dev(
    a2a_client: AsyncClient, monkeypatch: pytest.MonkeyPatch
) -> None:
    """A presented-but-forged token is rejected even in header-trust mode."""
    monkeypatch.setenv("ROBOCO_AGENT_AUTH_SECRET", _SECRET)
    monkeypatch.delenv("ROBOCO_AGENT_AUTH_REQUIRED", raising=False)
    r = await a2a_client.get(
        "/api/a2a/tasks/some-task/subscribe",
        headers={
            "X-Agent-ID": _AGENT_ID,
            "X-Agent-Role": "developer",
            "X-Agent-Token": "forged-not-a-real-hmac",
        },
    )
    assert r.status_code == _HTTP_401


@pytest.mark.asyncio
async def test_subscribe_accepts_valid_token_then_404s_unknown_task(
    a2a_client: AsyncClient, monkeypatch: pytest.MonkeyPatch
) -> None:
    """A valid token passes the gate; the route then 404s on the initial
    task-existence check (no DB seeded). 404 (not 401) proves the gate let
    the request through."""
    monkeypatch.setenv("ROBOCO_AGENT_AUTH_SECRET", _SECRET)
    monkeypatch.setenv("ROBOCO_AGENT_AUTH_REQUIRED", "true")
    token = issue_agent_token(_AGENT_ID, "developer")
    # get_task returns None -> 404. Patch A2AService.get_task to return None
    # so the route doesn't need a real DB.
    monkeypatch.setattr(a2a_module.A2AService, "get_task", AsyncMock(return_value=None))
    r = await a2a_client.get(
        "/api/a2a/tasks/some-task/subscribe",
        headers={
            "X-Agent-ID": _AGENT_ID,
            "X-Agent-Role": "developer",
            "X-Agent-Token": token,
        },
    )
    assert r.status_code == _HTTP_404


# ---------------------------------------------------------------------------
# Session-per-query: structural + behavioral
# ---------------------------------------------------------------------------


def test_subscribe_route_does_not_hold_request_scoped_db() -> None:
    """The route must NOT depend on ``get_db`` — the request-scoped session
    would be held for the full SSE lifetime (up to 1 hour). Each poll opens
    its own short-lived session via ``get_session_factory``.
    """
    subscribe_route = cast(
        "APIRoute",
        next(
            r
            for r in a2a_router.routes
            if getattr(r, "path", "") == "/tasks/{task_id}/subscribe"
        ),
    )
    # Walk the route's dependency tree; get_db must not appear anywhere.
    deps = [subscribe_route.dependant]
    seen: set[int] = set()
    found_get_db = False
    while deps:
        d = deps.pop()
        if id(d) in seen:
            continue
        seen.add(id(d))
        if d.call is get_db:
            found_get_db = True
        deps.extend(d.dependencies)
    assert not found_get_db, (
        "subscribe_to_task still depends on get_db — the request-scoped "
        "session is held for the full SSE lifetime (pool-exhaustion vector)."
    )


@pytest.mark.asyncio
async def test_subscribe_opens_a_short_lived_session_per_poll(
    a2a_client: AsyncClient, monkeypatch: pytest.MonkeyPatch
) -> None:
    """Each poll iteration opens its own session and closes it before the next
    ``asyncio.sleep`` — never holding one connection across the full SSE
    lifetime. Asserts more than one session open (one per poll, not one for
    the lifetime)."""

    monkeypatch.setenv("ROBOCO_AGENT_AUTH_SECRET", _SECRET)
    monkeypatch.setenv("ROBOCO_AGENT_AUTH_REQUIRED", "true")
    token = issue_agent_token(_AGENT_ID, "developer")

    # Count session opens across the SSE lifetime.
    open_count = {"n": 0}

    def _factory() -> Any:
        open_count["n"] += 1

        class _Ctx:
            async def __aenter__(self) -> MagicMock:
                return MagicMock()

            async def __aexit__(self, *exc: object) -> None:
                return None

        return _Ctx()

    monkeypatch.setattr(a2a_module, "get_session_factory", lambda: _factory)

    # Non-terminal fake task so the loop keeps polling.
    fake_task = MagicMock()
    fake_task.status.state = "in_progress"
    fake_task.model_dump_json = MagicMock(return_value="{}")
    monkeypatch.setattr(
        a2a_module.A2AService, "get_task", AsyncMock(return_value=fake_task)
    )

    # No sleeping — drain the generator as fast as possible.
    monkeypatch.setattr(a2a_module.asyncio, "sleep", AsyncMock(return_value=None))

    # Disconnect after 3 polls so the stream terminates.
    disconnect_after = {"remaining": 3}

    async def _fake_is_disconnected() -> bool:
        if disconnect_after["remaining"] <= 0:
            return True
        disconnect_after["remaining"] -= 1
        return False

    # The route reads request.is_disconnected(); patch it on the request via
    # the Starlette request. We patch the Request.is_disconnected property.
    monkeypatch.setattr(
        "fastapi.Request.is_disconnected",
        lambda _self: _fake_is_disconnected(),
    )

    r = await a2a_client.get(
        "/api/a2a/tasks/some-task/subscribe",
        headers={
            "X-Agent-ID": _AGENT_ID,
            "X-Agent-Role": "developer",
            "X-Agent-Token": token,
        },
    )
    # Drain the SSE stream so the generator runs to completion.
    assert r.status_code == _HTTP_200
    # Consume the body (the SSE stream finishes once is_disconnected returns
    # True on the 4th check).
    _ = await r.aread()

    # 3 polls + 1 initial validation = 4 session opens (one per query, none
    # held across the lifetime). The key assertion: more than one session
    # was opened — proving the request-scoped session is gone.
    assert open_count["n"] > 1, (
        f"only {open_count['n']} session open(s) — the route is holding a "
        "single request-scoped session for the full SSE lifetime (pool "
        "exhaustion vector)."
    )