feat: add agent zones
This commit is contained in:
lussierj committed
1 parent
d8cd7ad058
commit
82309db79a
29 files changed
+2242
-255
No files matched your search
@@ -0,0 +1,283 @@
|
||||
"""Focused model, service, router, and architecture tests for Agent Zones v1."""
|
||||
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from uuid import UUID, uuid4
|
||||
|
||||
from fastapi import FastAPI
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from app.agent_commands.dependencies import require_agent_identity
|
||||
from app.agent_zones.dependencies import get_agent_zone_service
|
||||
from app.agent_zones.exceptions import AgentZoneValidationError
|
||||
from app.agent_zones.models import AgentCheckinRequest, AgentZoneCreateRequest
|
||||
from app.agent_zones.router import router
|
||||
from app.agent_zones.service import AgentZoneService
|
||||
from app.identity import get_current_user, required_api_permission
|
||||
|
||||
|
||||
SITE_A = UUID("00000000-0000-4000-8000-000000001601")
|
||||
SITE_B = UUID("00000000-0000-4000-8000-000000001602")
|
||||
ZONE_A = UUID("00000000-0000-4000-8000-000000001611")
|
||||
ZONE_B = UUID("00000000-0000-4000-8000-000000001612")
|
||||
MINER_A = UUID("00000000-0000-4000-8000-000000001621")
|
||||
APP_ROOT = Path(__file__).parents[1] / "app"
|
||||
|
||||
|
||||
def zone(zone_id=ZONE_A, site_id=SITE_A, *, enabled=True, name="Main Room"):
|
||||
return {
|
||||
"id": zone_id,
|
||||
"site_id": site_id,
|
||||
"name": name,
|
||||
"description": None,
|
||||
"is_enabled": enabled,
|
||||
"assigned_agent_id": None,
|
||||
"networks": [],
|
||||
}
|
||||
|
||||
|
||||
class FakeRepository:
|
||||
def __init__(self):
|
||||
self.zones = {ZONE_A: zone()}
|
||||
self.agent = None
|
||||
self.saved = []
|
||||
self.assign_result = {MINER_A: "assigned"}
|
||||
|
||||
def get_agent(self, agent_id):
|
||||
return self.agent if self.agent and self.agent["agent_id"] == agent_id else None
|
||||
|
||||
def get_zone(self, zone_id):
|
||||
return self.zones.get(zone_id)
|
||||
|
||||
def list_usable_zones(self, site_id):
|
||||
return [
|
||||
value for value in self.zones.values()
|
||||
if value["site_id"] == site_id and value["is_enabled"]
|
||||
]
|
||||
|
||||
def list_zones(self, site_id=None):
|
||||
return [
|
||||
value for value in self.zones.values()
|
||||
if site_id is None or value["site_id"] == site_id
|
||||
]
|
||||
|
||||
def create_zone(self, *, site_id, name, description, networks):
|
||||
created_id = uuid4()
|
||||
self.zones[created_id] = {
|
||||
**zone(created_id, site_id, name=name),
|
||||
"description": description,
|
||||
"networks": [
|
||||
{"id": uuid4(), "network": network} for network in networks
|
||||
],
|
||||
}
|
||||
return created_id
|
||||
|
||||
def save_checkin(self, **kwargs):
|
||||
self.saved.append(kwargs)
|
||||
self.agent = {
|
||||
"agent_id": kwargs["agent_id"],
|
||||
"site_id": kwargs["site_id"],
|
||||
"agent_zone_id": kwargs["agent_zone_id"],
|
||||
"site_name": kwargs["site_name"],
|
||||
}
|
||||
|
||||
def get_agent_assignment(self, agent_id):
|
||||
if not self.agent or self.agent["agent_id"] != agent_id:
|
||||
return None
|
||||
zone_row = self.zones.get(self.agent.get("agent_zone_id"))
|
||||
return {
|
||||
**self.agent,
|
||||
"zone_enabled": zone_row["is_enabled"] if zone_row else None,
|
||||
}
|
||||
|
||||
def assign_miners_to_zone(self, miner_ids, agent_zone_id):
|
||||
self.assigned = (list(miner_ids), agent_zone_id)
|
||||
return {miner_id: self.assign_result.get(miner_id, "miner_not_found")
|
||||
for miner_id in miner_ids}
|
||||
|
||||
|
||||
class AgentZoneModelTests(unittest.TestCase):
|
||||
def test_names_and_networks_are_normalized(self):
|
||||
request = AgentZoneCreateRequest(
|
||||
site_id=SITE_A,
|
||||
name=" Main Room ",
|
||||
description=" Local VLAN ",
|
||||
networks=["10.30.0.0/24", "2001:db8::/64"],
|
||||
)
|
||||
self.assertEqual(request.name, "Main Room")
|
||||
self.assertEqual(request.description, "Local VLAN")
|
||||
self.assertEqual(request.networks, ["10.30.0.0/24", "2001:db8::/64"])
|
||||
|
||||
def test_invalid_host_network_duplicate_and_extra_fields_are_rejected(self):
|
||||
for networks in (["10.30.0.7/24"], ["10.30.0.0/24", "10.30.0.0/24"]):
|
||||
with self.subTest(networks=networks), self.assertRaises(Exception):
|
||||
AgentZoneCreateRequest(site_id=SITE_A, name="Zone", networks=networks)
|
||||
with self.assertRaises(Exception):
|
||||
AgentZoneCreateRequest(site_id=SITE_A, name="Zone", unknown=True)
|
||||
|
||||
|
||||
class AgentZoneServiceTests(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.repository = FakeRepository()
|
||||
self.service = AgentZoneService(self.repository)
|
||||
|
||||
def checkin(self, **kwargs):
|
||||
return self.service.checkin(
|
||||
"agent-a",
|
||||
AgentCheckinRequest(agent_id="agent-a", site_name="North", **kwargs),
|
||||
)
|
||||
|
||||
def test_explicit_first_assignment_and_same_assignment_are_allowed(self):
|
||||
first = self.checkin(site_id=SITE_A, agent_zone_id=ZONE_A)
|
||||
second = self.checkin(site_id=SITE_A, agent_zone_id=ZONE_A)
|
||||
self.assertTrue(first.routing_available)
|
||||
self.assertEqual(second.agent_zone_id, ZONE_A)
|
||||
|
||||
def test_explicit_zone_site_mismatch_and_identity_spoof_are_rejected(self):
|
||||
with self.assertRaisesRegex(AgentZoneValidationError, "does not belong"):
|
||||
self.checkin(site_id=SITE_B, agent_zone_id=ZONE_A)
|
||||
with self.assertRaisesRegex(AgentZoneValidationError, "authenticated"):
|
||||
self.service.checkin(
|
||||
"agent-a", AgentCheckinRequest(agent_id="agent-b", site_id=SITE_A)
|
||||
)
|
||||
|
||||
def test_established_agent_cannot_change_zone_within_site(self):
|
||||
self.checkin(site_id=SITE_A, agent_zone_id=ZONE_A)
|
||||
self.repository.zones[ZONE_B] = zone(ZONE_B, SITE_A, name="Room 2")
|
||||
with self.assertRaisesRegex(AgentZoneValidationError, "different Agent Zone"):
|
||||
self.checkin(site_id=SITE_A, agent_zone_id=ZONE_B)
|
||||
self.assertEqual(self.repository.agent["agent_zone_id"], ZONE_A)
|
||||
|
||||
def test_established_agent_site_conflict_precedes_zone_conflict(self):
|
||||
self.checkin(site_id=SITE_A, agent_zone_id=ZONE_A)
|
||||
self.repository.zones[ZONE_B] = zone(ZONE_B, SITE_B, name="Remote")
|
||||
with self.assertRaisesRegex(AgentZoneValidationError, "cannot change Sites"):
|
||||
self.checkin(site_id=SITE_B, agent_zone_id=ZONE_B)
|
||||
self.assertEqual(self.repository.agent["agent_zone_id"], ZONE_A)
|
||||
|
||||
def test_omitted_zone_preserves_established_assignment_and_site(self):
|
||||
self.checkin(site_id=SITE_A, agent_zone_id=ZONE_A)
|
||||
result = self.service.checkin(
|
||||
"agent-a", AgentCheckinRequest(agent_id="agent-a", site_name="North")
|
||||
)
|
||||
self.assertEqual(result.site_id, SITE_A)
|
||||
self.assertEqual(result.agent_zone_id, ZONE_A)
|
||||
|
||||
def test_transitional_resolution_requires_exactly_one_enabled_zone(self):
|
||||
single = self.checkin(site_id=SITE_A)
|
||||
self.assertEqual(single.agent_zone_id, ZONE_A)
|
||||
self.repository.agent = None
|
||||
self.repository.zones = {}
|
||||
zero = self.checkin(site_id=SITE_A)
|
||||
self.assertEqual(zero.routing_reason, "agent_zone_unavailable")
|
||||
self.repository.agent = None
|
||||
self.repository.zones = {
|
||||
ZONE_A: zone(),
|
||||
ZONE_B: zone(ZONE_B, SITE_A, name="Room 2"),
|
||||
}
|
||||
many = self.checkin(site_id=SITE_A)
|
||||
self.assertEqual(many.routing_reason, "agent_zone_ambiguous")
|
||||
self.assertIsNone(many.agent_zone_id)
|
||||
|
||||
def test_disabled_zone_is_not_selected_and_existing_assignment_is_factual(self):
|
||||
self.repository.zones[ZONE_A]["is_enabled"] = False
|
||||
unassigned = self.checkin(site_id=SITE_A)
|
||||
self.assertEqual(unassigned.routing_reason, "agent_zone_unavailable")
|
||||
self.repository.agent = {
|
||||
"agent_id": "agent-a", "site_id": SITE_A,
|
||||
"agent_zone_id": ZONE_A, "site_name": "North",
|
||||
}
|
||||
existing = self.service.checkin(
|
||||
"agent-a", AgentCheckinRequest(agent_id="agent-a")
|
||||
)
|
||||
self.assertEqual(existing.routing_reason, "agent_zone_disabled")
|
||||
self.assertEqual(existing.agent_zone_id, ZONE_A)
|
||||
|
||||
def test_payload_is_sanitized_before_persistence(self):
|
||||
request = AgentCheckinRequest.model_validate({
|
||||
"agent_id": "agent-a", "site_id": str(SITE_A),
|
||||
"api3_password": "secret", "nested": {"access_token": "secret"},
|
||||
"control_available": True,
|
||||
})
|
||||
self.service.checkin("agent-a", request)
|
||||
payload = self.repository.saved[-1]["payload"]
|
||||
self.assertNotIn("password", str(payload).lower())
|
||||
self.assertNotIn("token", str(payload).lower())
|
||||
self.assertTrue(payload["control_available"])
|
||||
|
||||
def test_miner_zone_assignment_surfaces_conflicts_without_rehoming(self):
|
||||
self.checkin(site_id=SITE_A, agent_zone_id=ZONE_A)
|
||||
self.repository.assign_result[MINER_A] = "agent_zone_conflict"
|
||||
result = self.service.associate_miners_from_agent("agent-a", [MINER_A])
|
||||
self.assertEqual(result.outcomes[MINER_A], "agent_zone_conflict")
|
||||
self.assertEqual(self.repository.assigned, ([MINER_A], ZONE_A))
|
||||
|
||||
def test_routing_context_rejects_stale_recorded_zone(self):
|
||||
self.checkin(site_id=SITE_A, agent_zone_id=ZONE_A)
|
||||
result = self.service.routing_context("agent-a", ZONE_B)
|
||||
self.assertEqual(result.routing_reason, "agent_zone_conflict")
|
||||
self.assertIsNone(result.site_id)
|
||||
|
||||
|
||||
class AgentZoneRouterTests(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.repository = FakeRepository()
|
||||
self.service = AgentZoneService(self.repository)
|
||||
self.app = FastAPI()
|
||||
self.app.include_router(router)
|
||||
self.app.dependency_overrides[get_agent_zone_service] = lambda: self.service
|
||||
self.app.dependency_overrides[require_agent_identity] = lambda: "agent-a"
|
||||
self.client = TestClient(self.app)
|
||||
|
||||
def tearDown(self):
|
||||
self.client.close()
|
||||
|
||||
def authenticate(self, permissions):
|
||||
self.app.dependency_overrides[get_current_user] = lambda: {
|
||||
"id": str(uuid4()), "username": "operator", "role": "Custom",
|
||||
"permissions": list(permissions),
|
||||
}
|
||||
|
||||
def test_list_and_create_use_existing_agent_permissions(self):
|
||||
self.assertEqual(self.client.get("/api/agent-zones").status_code, 401)
|
||||
self.authenticate(("agents.view",))
|
||||
self.assertEqual(self.client.get("/api/agent-zones").status_code, 200)
|
||||
self.assertEqual(
|
||||
self.client.post(
|
||||
"/api/agent-zones", json={"site_id": str(SITE_A), "name": "Room 2"}
|
||||
).status_code,
|
||||
403,
|
||||
)
|
||||
self.authenticate(("agents.manage",))
|
||||
response = self.client.post(
|
||||
"/api/agent-zones",
|
||||
json={"site_id": str(SITE_A), "name": "Room 2", "networks": ["10.40.0.0/24"]},
|
||||
)
|
||||
self.assertEqual(response.status_code, 201)
|
||||
self.assertEqual(response.json()["data"]["name"], "Room 2")
|
||||
|
||||
def test_checkin_uses_authenticated_agent_identity(self):
|
||||
response = self.client.post(
|
||||
"/api/agent/checkin",
|
||||
json={"agent_id": "agent-a", "site_id": str(SITE_A),
|
||||
"agent_zone_id": str(ZONE_A)},
|
||||
)
|
||||
self.assertEqual(response.status_code, 200)
|
||||
self.assertEqual(response.json()["agent_id"], "agent-a")
|
||||
self.assertNotIn("secret", response.text.lower())
|
||||
|
||||
def test_permission_and_architecture_audits(self):
|
||||
self.assertEqual(required_api_permission("GET", "/api/agent-zones"), "agents.view")
|
||||
self.assertEqual(required_api_permission("POST", "/api/agent-zones"), "agents.manage")
|
||||
router_source = (APP_ROOT / "agent_zones" / "router.py").read_text()
|
||||
service_source = (APP_ROOT / "agent_zones" / "service.py").read_text()
|
||||
repository_source = (APP_ROOT / "agent_zones" / "repository.py").read_text()
|
||||
self.assertNotIn("psycopg", router_source)
|
||||
self.assertNotIn("SELECT ", router_source)
|
||||
self.assertNotIn("fastapi", service_source.lower())
|
||||
self.assertNotIn("fastapi", repository_source.lower())
|
||||
self.assertNotIn('@app.post("/api/agent/checkin")', (APP_ROOT / "main.py").read_text())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in new issue
Block a user