320 lines
13 KiB
Python
320 lines
13 KiB
Python
"""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 get_agent_configuration(self, agent_id):
|
|
assignment = self.get_agent_assignment(agent_id)
|
|
if not assignment:
|
|
return None
|
|
zone_row = self.zones.get(assignment.get("agent_zone_id"))
|
|
return {
|
|
"agent_id": agent_id,
|
|
"site_id": assignment.get("site_id"),
|
|
"agent_zone_id": assignment.get("agent_zone_id"),
|
|
"site_name": assignment.get("site_name"),
|
|
"agent_zone_name": zone_row["name"] if zone_row else None,
|
|
"agent_zone_enabled": zone_row["is_enabled"] if zone_row else None,
|
|
"networks": [
|
|
item["network"] for item in (zone_row or {}).get("networks", [])
|
|
],
|
|
}
|
|
|
|
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)
|
|
|
|
def test_configuration_uses_only_authenticated_agent_zone_networks(self):
|
|
self.repository.zones[ZONE_A]["networks"] = [
|
|
{"id": uuid4(), "network": "10.30.0.0/24"}
|
|
]
|
|
self.repository.zones[ZONE_B] = zone(ZONE_B, SITE_B, name="Remote")
|
|
self.repository.zones[ZONE_B]["networks"] = [
|
|
{"id": uuid4(), "network": "10.80.0.0/24"}
|
|
]
|
|
self.checkin(site_id=SITE_A, agent_zone_id=ZONE_A)
|
|
|
|
configuration = self.service.configuration("agent-a")
|
|
|
|
self.assertEqual(configuration["agent_zone_id"], ZONE_A)
|
|
self.assertEqual(
|
|
[item["cidr"] for item in configuration["ip_ranges"]],
|
|
["10.30.0.0/24"],
|
|
)
|
|
self.assertNotIn("10.80.0.0/24", str(configuration))
|
|
|
|
|
|
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()
|