90 lines
3.9 KiB
Python
90 lines
3.9 KiB
Python
# tests/test_demo_dataplane_wiring.py
|
|
from __future__ import annotations
|
|
|
|
import unittest
|
|
from datetime import UTC, datetime, timedelta
|
|
from uuid import uuid4
|
|
|
|
from services.control_api.auth.models import AuthenticatedPrincipal, PrincipalKind
|
|
from services.runtime.dataplane import CompositeAuthenticator, DemoAwareRateLimiter
|
|
from services.runtime.demo_token import TOKEN_PREFIX
|
|
|
|
|
|
def _principal(subject: str) -> AuthenticatedPrincipal:
|
|
now = datetime.now(UTC)
|
|
return AuthenticatedPrincipal(
|
|
subject=subject,
|
|
issuer="https://fake.oidc.test",
|
|
audience=("itworx-mcp-hub",),
|
|
kind=PrincipalKind.MACHINE_CLIENT,
|
|
client_id=uuid4(),
|
|
expires_at=now + timedelta(minutes=5),
|
|
issued_at=now,
|
|
)
|
|
|
|
|
|
class CompositeAuthenticatorTests(unittest.TestCase):
|
|
def test_demo_shaped_bearer_routes_to_demo_authenticator(self) -> None:
|
|
demo_calls: list[str | None] = []
|
|
oidc_calls: list[str | None] = []
|
|
composite = CompositeAuthenticator(
|
|
demo=lambda header: (demo_calls.append(header), _principal("demo"))[1],
|
|
oidc=lambda header: (oidc_calls.append(header), _principal("real"))[1],
|
|
)
|
|
header = f"Bearer {TOKEN_PREFIX}abc.def"
|
|
result = composite(header)
|
|
self.assertEqual(result.subject, "demo")
|
|
self.assertEqual(demo_calls, [header])
|
|
self.assertEqual(oidc_calls, [])
|
|
|
|
def test_non_demo_bearer_routes_to_oidc_authenticator(self) -> None:
|
|
demo_calls: list[str | None] = []
|
|
oidc_calls: list[str | None] = []
|
|
composite = CompositeAuthenticator(
|
|
demo=lambda header: (demo_calls.append(header), _principal("demo"))[1],
|
|
oidc=lambda header: (oidc_calls.append(header), _principal("real"))[1],
|
|
)
|
|
header = "Bearer eyJhbGciOiJSUzI1NiJ9.fake.jwt"
|
|
result = composite(header)
|
|
self.assertEqual(result.subject, "real")
|
|
self.assertEqual(oidc_calls, [header])
|
|
self.assertEqual(demo_calls, [])
|
|
|
|
def test_none_header_routes_to_oidc_authenticator(self) -> None:
|
|
oidc_calls: list[str | None] = []
|
|
composite = CompositeAuthenticator(
|
|
demo=lambda header: _principal("demo"),
|
|
oidc=lambda header: (oidc_calls.append(header), _principal("real"))[1],
|
|
)
|
|
composite(None)
|
|
self.assertEqual(oidc_calls, [None])
|
|
|
|
|
|
class DemoAwareRateLimiterTests(unittest.TestCase):
|
|
def test_demo_client_gets_tighter_ceiling(self) -> None:
|
|
demo_client_id = uuid4()
|
|
limiter = DemoAwareRateLimiter(
|
|
default_limit=60, demo_limit=2, demo_client_id=demo_client_id
|
|
)
|
|
self.assertTrue(limiter.allow(client_id=demo_client_id, project="mobilityops", tool="t"))
|
|
self.assertTrue(limiter.allow(client_id=demo_client_id, project="mobilityops", tool="t"))
|
|
self.assertFalse(limiter.allow(client_id=demo_client_id, project="mobilityops", tool="t"))
|
|
|
|
def test_other_clients_keep_default_ceiling(self) -> None:
|
|
demo_client_id = uuid4()
|
|
other_client_id = uuid4()
|
|
limiter = DemoAwareRateLimiter(default_limit=3, demo_limit=1, demo_client_id=demo_client_id)
|
|
for _ in range(3):
|
|
self.assertTrue(limiter.allow(client_id=other_client_id, project="ludarium", tool="t"))
|
|
self.assertFalse(limiter.allow(client_id=other_client_id, project="ludarium", tool="t"))
|
|
|
|
def test_demo_and_other_clients_do_not_share_a_bucket(self) -> None:
|
|
demo_client_id = uuid4()
|
|
other_client_id = uuid4()
|
|
limiter = DemoAwareRateLimiter(
|
|
default_limit=60, demo_limit=1, demo_client_id=demo_client_id
|
|
)
|
|
self.assertTrue(limiter.allow(client_id=demo_client_id, project="mobilityops", tool="t"))
|
|
self.assertFalse(limiter.allow(client_id=demo_client_id, project="mobilityops", tool="t"))
|
|
self.assertTrue(limiter.allow(client_id=other_client_id, project="mobilityops", tool="t"))
|