139 lines
3.8 KiB
Python
139 lines
3.8 KiB
Python
from __future__ import annotations
|
|
|
|
import pytest
|
|
|
|
from src.services import awooop_approval_token as approval_token
|
|
|
|
|
|
class FakeRedis:
|
|
def __init__(self) -> None:
|
|
self.values: dict[str, str] = {}
|
|
self.sets: dict[str, set[str]] = {}
|
|
self.expirations: dict[str, int] = {}
|
|
|
|
async def set(
|
|
self,
|
|
key: str,
|
|
value: str,
|
|
*,
|
|
nx: bool = False,
|
|
ex: int | None = None,
|
|
) -> bool:
|
|
if nx and key in self.values:
|
|
return False
|
|
self.values[key] = value
|
|
if ex is not None:
|
|
self.expirations[key] = ex
|
|
return True
|
|
|
|
async def sadd(self, key: str, member: str) -> int:
|
|
members = self.sets.setdefault(key, set())
|
|
before = len(members)
|
|
members.add(member)
|
|
return 1 if len(members) > before else 0
|
|
|
|
async def expire(self, key: str, ttl: int) -> bool:
|
|
self.expirations[key] = ttl
|
|
return True
|
|
|
|
async def scard(self, key: str) -> int:
|
|
return len(self.sets.get(key, set()))
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def stable_hmac(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
monkeypatch.setattr(approval_token, "_get_hmac_key", lambda: b"test-hmac-key")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_record_approval_uses_shared_redis_pool(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
redis = FakeRedis()
|
|
monkeypatch.setattr(approval_token, "get_redis", lambda: redis)
|
|
|
|
token = approval_token.issue_approval_token(
|
|
project_id="awoooi",
|
|
run_id="run-1",
|
|
tool_name="operator_console_approve",
|
|
approver_id="telegram:123456",
|
|
)
|
|
|
|
count = await approval_token.record_approval(
|
|
project_id="awoooi",
|
|
run_id="run-1",
|
|
tool_name="operator_console_approve",
|
|
approver_id="telegram:123456",
|
|
token=token,
|
|
)
|
|
|
|
assert count == 1
|
|
assert any(key.startswith("awooop_appr:jti:") for key in redis.values)
|
|
assert redis.sets[
|
|
"awooop_appr:sigs:awoooi:run-1:operator_console_approve"
|
|
] == {"telegram:123456"}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_record_approval_rejects_token_replay(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
redis = FakeRedis()
|
|
monkeypatch.setattr(approval_token, "get_redis", lambda: redis)
|
|
|
|
token = approval_token.issue_approval_token(
|
|
project_id="awoooi",
|
|
run_id="run-1",
|
|
tool_name="operator_console_approve",
|
|
approver_id="telegram:123456",
|
|
)
|
|
|
|
await approval_token.record_approval(
|
|
project_id="awoooi",
|
|
run_id="run-1",
|
|
tool_name="operator_console_approve",
|
|
approver_id="telegram:123456",
|
|
token=token,
|
|
)
|
|
|
|
with pytest.raises(approval_token.TokenReplayError):
|
|
await approval_token.record_approval(
|
|
project_id="awoooi",
|
|
run_id="run-1",
|
|
tool_name="operator_console_approve",
|
|
approver_id="telegram:123456",
|
|
token=token,
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_check_approval_quorum_uses_shared_redis_pool(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
redis = FakeRedis()
|
|
redis.sets["awooop_appr:sigs:awoooi:run-1:tool"] = {"operator-a"}
|
|
monkeypatch.setattr(approval_token, "get_redis", lambda: redis)
|
|
|
|
assert await approval_token.check_approval_quorum(
|
|
project_id="awoooi",
|
|
run_id="run-1",
|
|
tool_name="tool",
|
|
required_count=1,
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_check_approval_quorum_rejects_insufficient_count(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
redis = FakeRedis()
|
|
monkeypatch.setattr(approval_token, "get_redis", lambda: redis)
|
|
|
|
with pytest.raises(approval_token.QuorumNotMetError):
|
|
await approval_token.check_approval_quorum(
|
|
project_id="awoooi",
|
|
run_id="run-1",
|
|
tool_name="tool",
|
|
required_count=1,
|
|
)
|