192 lines
8.4 KiB
Python
192 lines
8.4 KiB
Python
"""Client model, clients file and hot-reloading registry."""
|
|
|
|
import json
|
|
import os
|
|
import stat
|
|
import tempfile
|
|
import time
|
|
import unittest
|
|
from datetime import date, timedelta
|
|
from pathlib import Path
|
|
|
|
from mcp_calculator.clients import (
|
|
Client,
|
|
ClientConfigError,
|
|
ClientFile,
|
|
ClientRegistry,
|
|
clients_from_env,
|
|
hash_token,
|
|
new_token,
|
|
)
|
|
|
|
HASH = "a" * 64
|
|
|
|
|
|
class ClientModelTest(unittest.TestCase):
|
|
def test_token_client(self):
|
|
c = Client.from_dict("team-a", {"token_sha256": HASH, "max_concurrent": 3, "allowed_networks": ["10.0.0.0/8"]})
|
|
self.assertEqual(c.status(), "active")
|
|
self.assertEqual(c.max_concurrent, 3)
|
|
self.assertTrue(c.allows_ip("10.1.2.3"))
|
|
self.assertFalse(c.allows_ip("192.168.1.1"))
|
|
self.assertFalse(c.allows_ip(None))
|
|
|
|
def test_without_networks_every_ip_is_allowed(self):
|
|
self.assertTrue(Client.from_dict("a", {"token_sha256": HASH}).allows_ip("203.0.113.9"))
|
|
|
|
def test_status(self):
|
|
today = date(2027, 1, 10)
|
|
self.assertEqual(Client.from_dict("a", {"token_sha256": HASH, "enabled": False}).status(today), "disabled")
|
|
self.assertEqual(Client.from_dict("a", {"token_sha256": HASH, "expires": "2027-01-09"}).status(today), "expired")
|
|
self.assertEqual(Client.from_dict("a", {"token_sha256": HASH, "expires": "2027-01-10"}).status(today), "active")
|
|
|
|
def test_network_client(self):
|
|
c = Client.from_dict("batch", {"auth": "network", "allowed_networks": ["10.131.2.15/32"]})
|
|
self.assertEqual(c.auth, "network")
|
|
self.assertTrue(c.allows_ip("10.131.2.15"))
|
|
|
|
def test_invalid_entries(self):
|
|
cases = {
|
|
"bad name": ("Team A", {"token_sha256": HASH}),
|
|
"missing hash": ("a", {}),
|
|
"short hash": ("a", {"token_sha256": "abc"}),
|
|
"unknown field": ("a", {"token_sha256": HASH, "admin": True}),
|
|
"bad auth": ("a", {"auth": "magic", "token_sha256": HASH}),
|
|
"network without networks": ("a", {"auth": "network"}),
|
|
"network with token": ("a", {"auth": "network", "token_sha256": HASH, "allowed_networks": ["10.0.0.1"]}),
|
|
"bad network": ("a", {"token_sha256": HASH, "allowed_networks": ["10.0.0.300/8"]}),
|
|
"bad date": ("a", {"token_sha256": HASH, "expires": "tomorrow"}),
|
|
"negative limit": ("a", {"token_sha256": HASH, "max_concurrent": -1}),
|
|
"enabled as string": ("a", {"token_sha256": HASH, "enabled": "yes"}),
|
|
}
|
|
for label, (name, data) in cases.items():
|
|
with self.subTest(label), self.assertRaises(ClientConfigError):
|
|
Client.from_dict(name, data)
|
|
|
|
def test_with_changes(self):
|
|
c = Client.from_dict("a", {"token_sha256": HASH})
|
|
changed = c.with_changes({"enabled": False, "max_concurrent": 5, "expires": "2030-01-01", "allowed_networks": ["10.0.0.0/8"]})
|
|
self.assertFalse(changed.enabled)
|
|
self.assertEqual(changed.max_concurrent, 5)
|
|
self.assertEqual(changed.expires, date(2030, 1, 1))
|
|
self.assertIsNotNone(changed.updated)
|
|
with self.assertRaisesRegex(ClientConfigError, "cannot change"):
|
|
c.with_changes({"token_sha256": "b" * 64})
|
|
|
|
def test_public_view_hides_hash(self):
|
|
view = Client.from_dict("a", {"token_sha256": HASH, "token_hint": "mcpc_abcd"}).public_view()
|
|
self.assertNotIn("token_sha256", view)
|
|
self.assertEqual(view["token_hint"], "mcpc_abcd")
|
|
self.assertEqual(view["status"], "active")
|
|
|
|
def test_tokens(self):
|
|
token = new_token()
|
|
self.assertTrue(token.startswith("mcpc_"))
|
|
self.assertGreaterEqual(len(token), 40)
|
|
self.assertNotEqual(token, new_token())
|
|
self.assertEqual(len(hash_token(token)), 64)
|
|
|
|
def test_clients_from_env(self):
|
|
clients = clients_from_env(["svc:0123456789abcdef"])
|
|
self.assertEqual(clients["svc"].token_sha256, hash_token("0123456789abcdef"))
|
|
self.assertEqual(clients["svc"].source, "env")
|
|
|
|
|
|
class ClientFileTest(unittest.TestCase):
|
|
def setUp(self):
|
|
self.tmp = tempfile.TemporaryDirectory()
|
|
self.addCleanup(self.tmp.cleanup)
|
|
self.path = Path(self.tmp.name) / "data" / "clients.json"
|
|
self.file = ClientFile(self.path)
|
|
|
|
def test_missing_file_is_empty(self):
|
|
self.assertEqual(self.file.read(), {})
|
|
|
|
def test_update_roundtrip(self):
|
|
with self.file.update() as clients:
|
|
clients["a"] = Client.from_dict("a", {"token_sha256": HASH, "note": "Team A"})
|
|
raw = json.loads(self.path.read_text())
|
|
self.assertEqual(raw["version"], 1)
|
|
self.assertEqual(raw["clients"]["a"]["note"], "Team A")
|
|
self.assertEqual(self.file.read()["a"].token_sha256, HASH)
|
|
self.assertEqual(stat.S_IMODE(os.stat(self.path).st_mode), 0o640)
|
|
self.assertEqual([p.name for p in self.path.parent.iterdir() if p.suffix == ".tmp"], [])
|
|
|
|
def test_failed_update_leaves_file_untouched(self):
|
|
with self.file.update() as clients:
|
|
clients["a"] = Client.from_dict("a", {"token_sha256": HASH})
|
|
before = self.path.read_text()
|
|
with self.assertRaises(RuntimeError), self.file.update() as clients:
|
|
clients.clear()
|
|
raise RuntimeError("abort")
|
|
self.assertEqual(self.path.read_text(), before)
|
|
|
|
def test_invalid_json(self):
|
|
self.path.parent.mkdir(parents=True)
|
|
self.path.write_text("{not json")
|
|
with self.assertRaisesRegex(ClientConfigError, "invalid JSON"):
|
|
self.file.read()
|
|
|
|
def test_duplicate_tokens_rejected(self):
|
|
self.path.parent.mkdir(parents=True)
|
|
self.path.write_text(json.dumps({"clients": {"a": {"token_sha256": HASH}, "b": {"token_sha256": HASH}}}))
|
|
with self.assertRaisesRegex(ClientConfigError, "same token"):
|
|
self.file.read()
|
|
|
|
|
|
class RegistryTest(unittest.TestCase):
|
|
def setUp(self):
|
|
self.tmp = tempfile.TemporaryDirectory()
|
|
self.addCleanup(self.tmp.cleanup)
|
|
self.file = ClientFile(Path(self.tmp.name) / "clients.json")
|
|
|
|
def _write(self, clients: dict) -> None:
|
|
with self.file.update() as current:
|
|
current.clear()
|
|
current.update({n: Client.from_dict(n, d) for n, d in clients.items()})
|
|
# make sure the fingerprint changes even on file systems with coarse timestamps
|
|
future = time.time() + 1 + len(clients)
|
|
os.utime(self.file.path, (future, future))
|
|
|
|
def test_hot_reload(self):
|
|
self._write({"a": {"token_sha256": HASH}})
|
|
registry = ClientRegistry(self.file, {}, reload_seconds=0.01)
|
|
self.assertEqual(set(registry.snapshot().clients), {"a"})
|
|
self._write({"a": {"token_sha256": HASH}, "b": {"token_sha256": "b" * 64}})
|
|
time.sleep(0.02)
|
|
self.assertEqual(set(registry.snapshot().clients), {"a", "b"})
|
|
self.assertIn("b" * 64, registry.snapshot().by_hash)
|
|
|
|
def test_invalid_change_keeps_previous_state(self):
|
|
self._write({"a": {"token_sha256": HASH}})
|
|
registry = ClientRegistry(self.file, {}, reload_seconds=0.01)
|
|
self.file.path.write_text("{broken")
|
|
os.utime(self.file.path, (time.time() + 5, time.time() + 5))
|
|
time.sleep(0.02)
|
|
with self.assertLogs("mcp_calculator.clients", "ERROR"):
|
|
snapshot = registry.snapshot()
|
|
self.assertEqual(set(snapshot.clients), {"a"})
|
|
|
|
def test_invalid_file_at_startup_fails(self):
|
|
self.file.path.write_text("{broken")
|
|
with self.assertRaises(ClientConfigError):
|
|
ClientRegistry(self.file, {})
|
|
|
|
def test_static_and_file_clients_are_merged(self):
|
|
self._write({"a": {"token_sha256": HASH}})
|
|
registry = ClientRegistry(self.file, clients_from_env(["svc:0123456789abcdef"]))
|
|
self.assertEqual(set(registry.snapshot().clients), {"a", "svc"})
|
|
|
|
def test_name_clash_between_sources_fails(self):
|
|
self._write({"svc": {"token_sha256": HASH}})
|
|
with self.assertRaisesRegex(ClientConfigError, "both"):
|
|
ClientRegistry(self.file, clients_from_env(["svc:0123456789abcdef"]))
|
|
|
|
def test_disabled_without_sources(self):
|
|
self.assertFalse(ClientRegistry(None, {}).enabled)
|
|
self.assertTrue(ClientRegistry(self.file, {}).enabled)
|
|
|
|
def test_expiry_is_evaluated_at_request_time(self):
|
|
yesterday = (date.today() - timedelta(days=1)).isoformat()
|
|
self._write({"old": {"token_sha256": HASH, "expires": yesterday}})
|
|
self.assertEqual(ClientRegistry(self.file, {}).snapshot().clients["old"].status(), "expired")
|