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(), )