Enterprise-MCP-Service-Bus/RI/tests/e2e/conftest.py

"""Shared e2e fixtures: a real OPA process, a real JWKS HTTP endpoint + RS256
token signing, and helpers to run the Gateway/Sales servers in-process
against them. No Keycloak available in this environment (no Java) -- see
RI/README.md for exactly what that does and does not validate.
"""

from __future__ import annotations

import asyncio
import json
import os
import shutil
import socket
import subprocess
import threading
import time
from collections.abc import AsyncIterator, Awaitable, Callable, Iterator
from http.server import BaseHTTPRequestHandler, HTTPServer
from pathlib import Path

import jwt
import pytest
import uvicorn
from cryptography.hazmat.primitives.asymmetric import rsa
from jwt.algorithms import RSAAlgorithm
from mcp.client.session import ClientSession
from mcp.client.streamable_http import (  # type: ignore[attr-defined]
    create_mcp_http_client,
    streamable_http_client,
)
from mcp.shared.exceptions import MCPError

REPO_ROOT = Path(__file__).parents[2]
POLICIES_DIR = REPO_ROOT / "config" / "policies"

ISSUER = "https://keycloak.example/realms/emcp"
AUDIENCE = "emcp-gateway"
KID = "test-key-1"

_current_gateway_port: int | None = None
"""Set by `governed_stack` (or a test managing its own stack, via
`set_current_gateway_port`) for the duration of one test. A *fresh* port
per test -- rather than a fixed constant -- avoids TIME_WAIT/bind
contention between one test's `RunningServer.stop()` and the next test's
bind on the same port, which showed up as intermittent "SSE stream ended
without a response" failures once enough e2e tests accumulated to make
that race likely within a single pytest run."""


def set_current_gateway_port(port: int | None) -> None:
    global _current_gateway_port
    _current_gateway_port = port


def _gateway_url() -> str:
    if _current_gateway_port is None:
        raise RuntimeError(
            "No gateway is currently running for this test (see governed_stack fixture "
            "or set_current_gateway_port)"
        )
    return f"http://127.0.0.1:{_current_gateway_port}/mcp"


def gateway_url() -> str:
    """Public accessor for `tests/adversarial/`: tests that need to send a raw
    (non-`ClientSession`) HTTP request against the currently running
    `governed_stack` gateway -- e.g. a header/body mismatch that must be
    rejected before any MCP session handshake."""
    return _gateway_url()


def free_port() -> int:
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
        s.bind(("127.0.0.1", 0))
        return int(s.getsockname()[1])


# --- Real OPA process (EP-03) --------------------------------------------


@pytest.fixture(scope="session")
def opa_url() -> Iterator[str]:
    opa_binary = shutil.which(os.environ.get("OPA_BINARY", "opa"))
    if opa_binary is None:
        pytest.skip(
            "opa binary not found on PATH (set OPA_BINARY or install it -- see RI/README.md)"
        )

    port = free_port()
    process = subprocess.Popen(
        [opa_binary, "run", "--server", "--addr", f"127.0.0.1:{port}", str(POLICIES_DIR)],
        stdout=subprocess.DEVNULL,
        stderr=subprocess.DEVNULL,
    )
    base_url = f"http://127.0.0.1:{port}"
    try:
        import httpx

        for _ in range(100):
            try:
                httpx.get(f"{base_url}/health", timeout=1)
                break
            except httpx.HTTPError:
                time.sleep(0.05)
        else:
            raise RuntimeError("opa server did not become healthy in time")
        yield base_url
    finally:
        process.terminate()
        process.wait(timeout=5)


# --- Real JWKS HTTP endpoint + RS256 signing (EP-01-T03) ------------------


class _JWKSHandler(BaseHTTPRequestHandler):
    jwks_document: dict[str, object] = {}

    def do_GET(self) -> None:  # noqa: N802
        body = json.dumps(self.jwks_document).encode("utf-8")
        self.send_response(200)
        self.send_header("Content-Type", "application/json")
        self.end_headers()
        self.wfile.write(body)

    def log_message(self, format: str, *args: object) -> None:  # noqa: A002
        pass


@pytest.fixture(scope="session")
def rsa_keypair() -> tuple[rsa.RSAPrivateKey, rsa.RSAPublicKey]:
    private_key = rsa.generate_private_key(public_exponent=65537, key_size=2048)
    return private_key, private_key.public_key()


@pytest.fixture(scope="session")
def jwks_url(rsa_keypair: tuple[rsa.RSAPrivateKey, rsa.RSAPublicKey]) -> Iterator[str]:
    _, public_key = rsa_keypair
    jwk = RSAAlgorithm.to_jwk(public_key, as_dict=True)
    jwk.update({"kid": KID, "use": "sig", "alg": "RS256"})
    _JWKSHandler.jwks_document = {"keys": [jwk]}

    server = HTTPServer(("127.0.0.1", 0), _JWKSHandler)
    thread = threading.Thread(target=server.serve_forever, daemon=True)
    thread.start()
    try:
        yield f"http://127.0.0.1:{server.server_port}/jwks"
    finally:
        server.shutdown()
        thread.join(timeout=5)


@pytest.fixture
def sign_token(
    rsa_keypair: tuple[rsa.RSAPrivateKey, rsa.RSAPublicKey],
) -> Callable[..., str]:
    private_key, _ = rsa_keypair

    def _sign(
        client_id: str,
        *,
        issuer: str = ISSUER,
        audience: str = AUDIENCE,
        expires_in_seconds: int = 300,
        signing_key: rsa.RSAPrivateKey | None = None,
    ) -> str:
        now = int(time.time())
        claims = {
            "iss": issuer,
            "aud": audience,
            "sub": f"service-account-{client_id}",
            "azp": client_id,
            "iat": now,
            "exp": now + expires_in_seconds,
        }
        return jwt.encode(
            claims, signing_key or private_key, algorithm="RS256", headers={"kid": KID}
        )

    return _sign


# --- Running Gateway + Sales servers, wired to the fixtures above --------


class RunningServer:
    def __init__(self, config: uvicorn.Config) -> None:
        self.server = uvicorn.Server(config)
        self._task: asyncio.Task[None] | None = None

    async def start(self) -> None:
        self._task = asyncio.create_task(self.server.serve())
        while not self.server.started:
            await asyncio.sleep(0.01)
        # See test_walking_skeleton.py's original note: give the ASGI app's
        # own lifespan (session manager task group) one extra tick.
        await asyncio.sleep(0.05)

    async def stop(self) -> None:
        self.server.should_exit = True
        assert self._task is not None
        await self._task


@pytest.fixture
async def governed_stack(
    monkeypatch: pytest.MonkeyPatch, opa_url: str, jwks_url: str
) -> AsyncIterator[None]:
    """Real OPA + real JWKS-backed OIDC auth, wired into fresh Gateway/Sales
    instances for one test (EP-01/EP-03/EP-04/EP-05 all real, only Keycloak
    itself is substituted by the JWKS server -- see RI/README.md).

    Deliberately Sales-only, not Sales+Finance: most tests using this fixture
    never reach a live backend call for a Finance tool (denied earlier, at
    entitlement or REQUIRE_APPROVAL), and this file's own flakiness note
    documents that more in-process ASGI server pairs per test measurably
    increases the known e2e flakiness rate. `governed_stack_with_finance`
    below is the (currently two) tests that actually need a live Finance
    backend -- keeping the blast radius of that cost limited to just those.
    """
    gateway_port = free_port()
    sales_port = free_port()

    monkeypatch.setenv("SALES_DOMAIN_URL", f"http://127.0.0.1:{sales_port}/mcp")
    monkeypatch.setenv("OPA_URL", opa_url)
    monkeypatch.setenv("OIDC_JWKS_URL", jwks_url)
    monkeypatch.setenv("OIDC_ISSUER", ISSUER)
    monkeypatch.setenv("OIDC_AUDIENCE", AUDIENCE)
    # EP-07: the Gateway's outbound credential to the sales-domain backend --
    # deliberately a different value/shape than any inbound client token.
    monkeypatch.setenv("SALES_DOMAIN_SERVICE_TOKEN", "test-sales-domain-backend-credential")

    import example_mcp_servers.sales_domain.server as sales_module

    sales_config = uvicorn.Config(
        sales_module.server.streamable_http_app(host="127.0.0.1"),
        host="127.0.0.1",
        port=sales_port,
        log_level="warning",
    )
    sales_server = RunningServer(sales_config)
    await sales_server.start()

    import emcp_bus.gateway.server as gateway_module

    auth_settings, token_verifier = gateway_module.build_auth()
    gateway_config = uvicorn.Config(
        gateway_module.build_app(auth_settings, token_verifier, host="127.0.0.1"),
        host="127.0.0.1",
        port=gateway_port,
        log_level="warning",
    )
    gateway_server = RunningServer(gateway_config)
    await gateway_server.start()
    set_current_gateway_port(gateway_port)

    try:
        yield
    finally:
        set_current_gateway_port(None)
        await gateway_server.stop()
        await sales_server.stop()


@pytest.fixture
async def governed_stack_with_finance(
    monkeypatch: pytest.MonkeyPatch, opa_url: str, jwks_url: str
) -> AsyncIterator[None]:
    """Same as `governed_stack`, plus a real Finance domain server (EP-06-T03) --
    for the handful of tests that need a live Finance backend to actually
    reach (e.g. proving the Fabric router dispatches Finance capabilities
    there, or a full approval-then-execute cycle against it)."""
    gateway_port = free_port()
    sales_port = free_port()
    finance_port = free_port()

    monkeypatch.setenv("SALES_DOMAIN_URL", f"http://127.0.0.1:{sales_port}/mcp")
    monkeypatch.setenv("FINANCE_DOMAIN_URL", f"http://127.0.0.1:{finance_port}/mcp")
    monkeypatch.setenv("OPA_URL", opa_url)
    monkeypatch.setenv("OIDC_JWKS_URL", jwks_url)
    monkeypatch.setenv("OIDC_ISSUER", ISSUER)
    monkeypatch.setenv("OIDC_AUDIENCE", AUDIENCE)
    monkeypatch.setenv("SALES_DOMAIN_SERVICE_TOKEN", "test-sales-domain-backend-credential")
    monkeypatch.setenv("FINANCE_DOMAIN_SERVICE_TOKEN", "test-finance-domain-backend-credential")

    import example_mcp_servers.sales_domain.server as sales_module

    sales_server = RunningServer(
        uvicorn.Config(
            sales_module.server.streamable_http_app(host="127.0.0.1"),
            host="127.0.0.1",
            port=sales_port,
            log_level="warning",
        )
    )
    await sales_server.start()

    import example_mcp_servers.finance_domain.server as finance_module

    finance_server = RunningServer(
        uvicorn.Config(
            finance_module.server.streamable_http_app(host="127.0.0.1"),
            host="127.0.0.1",
            port=finance_port,
            log_level="warning",
        )
    )
    await finance_server.start()

    import emcp_bus.gateway.server as gateway_module

    auth_settings, token_verifier = gateway_module.build_auth()
    gateway_server = RunningServer(
        uvicorn.Config(
            gateway_module.build_app(auth_settings, token_verifier, host="127.0.0.1"),
            host="127.0.0.1",
            port=gateway_port,
            log_level="warning",
        )
    )
    await gateway_server.start()
    set_current_gateway_port(gateway_port)

    try:
        yield
    finally:
        set_current_gateway_port(None)
        await gateway_server.stop()
        await finance_server.stop()
        await sales_server.stop()


async def run_against_gateway[ResultT](
    operation: Callable[[ClientSession], Awaitable[ResultT]],
    *,
    token: str | None = None,
) -> ResultT:
    """Connect to the gateway (optionally authenticated) and run `operation(session)`.

    Retries on `MCPError` ("SSE stream ended without a response" has been
    observed, rarely, under this in-process dual-uvicorn test setup under
    scheduler contention) -- a test-harness robustness concern, not a
    Gateway defect: production deployments run one process per service.
    """
    headers = {"Authorization": f"Bearer {token}"} if token else None
    url = _gateway_url()
    last_error: MCPError | None = None
    for attempt in range(3):
        if attempt:
            await asyncio.sleep(0.1 * attempt)
        try:
            async with (
                create_mcp_http_client(headers=headers) as http_client,
                streamable_http_client(url, http_client=http_client) as (read, write),
                ClientSession(read, write) as session,
            ):
                await session.initialize()
                return await operation(session)
        except MCPError as exc:
            last_error = exc
        except BaseExceptionGroup as group:
            # Each nested anyio TaskGroup (`streamable_http_client`'s, then
            # `ClientSession`'s) wraps the one below it in its own
            # ExceptionGroup, so a handler-raised MCPError (e.g. an auth
            # rejection during `initialize`) can arrive doubly-wrapped.
            # Unwrap fully so callers can `pytest.raises(MCPError)`
            # uniformly, same as for one raised directly.
            unwrapped = _unwrap_single_mcp_error(group)
            if unwrapped is not None:
                last_error = unwrapped
            else:
                raise
    assert last_error is not None
    raise last_error


def _unwrap_single_mcp_error(exc: BaseException) -> MCPError | None:
    """If `exc` is an MCPError, or a (possibly nested) ExceptionGroup wrapping
    exactly one, return it; otherwise `None`."""
    if isinstance(exc, MCPError):
        return exc
    if isinstance(exc, BaseExceptionGroup) and len(exc.exceptions) == 1:
        return _unwrap_single_mcp_error(exc.exceptions[0])
    return None