awoooi/apps/api/tests/test_mcp_gateway_gate5.py

from __future__ import annotations

import uuid

import pytest

from src.plugins.mcp import gateway as gateway_module
from src.plugins.mcp.gateway import (
    GateApprovalError,
    GateCheckResult,
    GateRateLimitError,
    GatewayContext,
    McpGateway,
)


class FakeRedis:
    def __init__(self, values: dict[str, str] | None = None) -> None:
        self.values = values or {}
        self.closed = False

    async def get(self, key: str) -> str | None:
        return self.values.get(key)

    async def aclose(self) -> None:
        self.closed = True

    async def eval(self, _script: str, _keys: int, key: str, _window: int) -> int:
        value = int(self.values.get(key, "0")) + 1
        self.values[key] = str(value)
        return value


@pytest.mark.asyncio
async def test_gate5_uses_shared_redis_pool_without_closing(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    run_id = uuid.uuid4()
    approval_key = f"mcp_approval:awoooi:openclaw-sre:k8s_apply:{run_id}"
    redis = FakeRedis({approval_key: "approved"})
    monkeypatch.setattr(gateway_module, "get_redis", lambda: redis)

    gate_result = GateCheckResult()
    await McpGateway(db=None)._gate5_approval(
        GatewayContext(
            project_id="awoooi",
            agent_id="openclaw-sre",
            tool_name="k8s_apply",
            run_id=run_id,
            is_shadow=False,
            required_scope="write",
        ),
        grant_row=None,
        gate_result=gate_result,
    )

    assert gate_result.gate5_approval is True
    assert redis.closed is False


@pytest.mark.asyncio
async def test_gate5_rejects_missing_approval(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    redis = FakeRedis()
    monkeypatch.setattr(gateway_module, "get_redis", lambda: redis)

    with pytest.raises(GateApprovalError):
        await McpGateway(db=None)._gate5_approval(
            GatewayContext(
                project_id="awoooi",
                agent_id="openclaw-sre",
                tool_name="k8s_apply",
                run_id=uuid.uuid4(),
                is_shadow=False,
                required_scope="write",
            ),
            grant_row=None,
            gate_result=GateCheckResult(),
        )


@pytest.mark.asyncio
async def test_gate5_read_scope_skips_redis(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    called = False

    def fail_if_called() -> FakeRedis:
        nonlocal called
        called = True
        return FakeRedis()

    monkeypatch.setattr(gateway_module, "get_redis", fail_if_called)
    gate_result = GateCheckResult()

    await McpGateway(db=None)._gate5_approval(
        GatewayContext(
            project_id="awoooi",
            agent_id="openclaw-sre",
            tool_name="k8s_get",
            is_shadow=False,
            required_scope="read",
        ),
        grant_row=None,
        gate_result=gate_result,
    )

    assert called is False
    assert gate_result.gate5_approval is True


@pytest.mark.asyncio
async def test_gateway_distributed_rate_limit_uses_redis_and_sets_receipt(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    redis = FakeRedis()
    monkeypatch.setattr(gateway_module, "get_redis", lambda: redis)
    gate_result = GateCheckResult()

    await McpGateway(db=None)._enforce_distributed_rate_limit(
        GatewayContext(
            project_id="awoooi",
            agent_id="openclaw-sre",
            tool_name="k8s_get",
            required_scope="read",
        ),
        gate_result,
    )

    assert gate_result.distributed_rate_limit is True
    assert redis.values[
        "mcp_rate_limit:awoooi:openclaw-sre:k8s_get:read"
    ] == "1"


@pytest.mark.asyncio
async def test_gateway_distributed_rate_limit_fails_closed_when_redis_is_down(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    def unavailable() -> FakeRedis:
        raise RuntimeError("redis unavailable")

    monkeypatch.setattr(gateway_module, "get_redis", unavailable)

    with pytest.raises(GateRateLimitError):
        await McpGateway(db=None)._enforce_distributed_rate_limit(
            GatewayContext(
                project_id="awoooi",
                agent_id="openclaw-sre",
                tool_name="k8s_apply",
                required_scope="write",
            ),
            GateCheckResult(),
        )