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,244 @@
|
||||
"""Real PostgreSQL integrity and routing tests for Agent Zones v1."""
|
||||
|
||||
import unittest
|
||||
from uuid import UUID
|
||||
|
||||
import psycopg
|
||||
from psycopg.errors import ForeignKeyViolation, UniqueViolation
|
||||
|
||||
from app.agent_zones.repository import AgentZoneRepository
|
||||
from app.config import settings
|
||||
from app.miner_control.repository import MinerControlRepository
|
||||
from app.miner_control.service import MinerControlService
|
||||
|
||||
|
||||
SITE_A = UUID("00000000-0000-4000-8000-000000001630")
|
||||
SITE_B = UUID("00000000-0000-4000-8000-000000001631")
|
||||
ZONE_A = UUID("00000000-0000-4000-8000-000000001632")
|
||||
ZONE_B = UUID("00000000-0000-4000-8000-000000001633")
|
||||
ZONE_C = UUID("00000000-0000-4000-8000-000000001636")
|
||||
MINER_UNASSIGNED = UUID("00000000-0000-4000-8000-000000001634")
|
||||
MINER_OTHER = UUID("00000000-0000-4000-8000-000000001635")
|
||||
AGENT_A = "agent-zone-postgres-a"
|
||||
|
||||
|
||||
def connect():
|
||||
return psycopg.connect(
|
||||
host=settings.postgres_host,
|
||||
port=settings.postgres_port,
|
||||
dbname=settings.postgres_db,
|
||||
user=settings.postgres_user,
|
||||
password=settings.postgres_password,
|
||||
connect_timeout=5,
|
||||
)
|
||||
|
||||
|
||||
class AgentZonePostgresTests(unittest.TestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
connection = connect()
|
||||
try:
|
||||
with connection.cursor() as cur:
|
||||
cur.execute(
|
||||
"INSERT INTO sites (id, name) VALUES (%s, 'Zone Test A'), "
|
||||
"(%s, 'Zone Test B') ON CONFLICT (id) DO NOTHING",
|
||||
(SITE_A, SITE_B),
|
||||
)
|
||||
cur.execute(
|
||||
"INSERT INTO agent_zones (id, site_id, name) VALUES "
|
||||
"(%s, %s, 'Main Room'), (%s, %s, 'Remote Room'), "
|
||||
"(%s, %s, 'Unstaffed Zone') "
|
||||
"ON CONFLICT (id) DO UPDATE SET is_enabled = true",
|
||||
(ZONE_A, SITE_A, ZONE_B, SITE_B, ZONE_C, SITE_A),
|
||||
)
|
||||
cur.execute(
|
||||
"INSERT INTO miners (id, site_id, asset_tag, current_ip, status, "
|
||||
"api3_supported, first_seen_at, last_seen_at) VALUES "
|
||||
"(%s, %s, 'ZONE-UNASSIGNED', '192.0.2.31', 'online', true, NULL, NULL), "
|
||||
"(%s, %s, 'ZONE-OTHER', '192.0.2.32', 'online', true, NULL, NULL) "
|
||||
"ON CONFLICT (id) DO UPDATE SET agent_zone_id = NULL",
|
||||
(MINER_UNASSIGNED, SITE_A, MINER_OTHER, SITE_B),
|
||||
)
|
||||
cur.execute(
|
||||
"UPDATE miners SET agent_zone_id = %s WHERE id = %s",
|
||||
(ZONE_B, MINER_OTHER),
|
||||
)
|
||||
cur.execute(
|
||||
"INSERT INTO minerops_site_agents ("
|
||||
"agent_id, site_id, agent_zone_id, site_name, last_payload) "
|
||||
"VALUES (%s, %s, %s, 'Zone Test A', '{\"control_available\":true}') "
|
||||
"ON CONFLICT (agent_id) DO UPDATE SET site_id = EXCLUDED.site_id, "
|
||||
"agent_zone_id = EXCLUDED.agent_zone_id, last_payload = EXCLUDED.last_payload",
|
||||
(AGENT_A, SITE_A, ZONE_A),
|
||||
)
|
||||
connection.commit()
|
||||
finally:
|
||||
connection.close()
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
connection = connect()
|
||||
try:
|
||||
with connection.cursor() as cur:
|
||||
cur.execute("DELETE FROM minerops_site_agents WHERE agent_id LIKE 'agent-zone-postgres-%'")
|
||||
cur.execute("DELETE FROM miners WHERE id IN (%s, %s)", (MINER_UNASSIGNED, MINER_OTHER))
|
||||
cur.execute("DELETE FROM agent_zone_network_ranges WHERE agent_zone_id IN (%s, %s, %s)", (ZONE_A, ZONE_B, ZONE_C))
|
||||
cur.execute("DELETE FROM agent_zones WHERE id IN (%s, %s, %s)", (ZONE_A, ZONE_B, ZONE_C))
|
||||
cur.execute("DELETE FROM sites WHERE id IN (%s, %s)", (SITE_A, SITE_B))
|
||||
connection.commit()
|
||||
finally:
|
||||
connection.close()
|
||||
|
||||
def setUp(self):
|
||||
connection = connect()
|
||||
try:
|
||||
with connection.cursor() as cur:
|
||||
cur.execute(
|
||||
"UPDATE agent_zones SET is_enabled = true "
|
||||
"WHERE id IN (%s, %s, %s)",
|
||||
(ZONE_A, ZONE_B, ZONE_C),
|
||||
)
|
||||
cur.execute(
|
||||
"UPDATE miners SET agent_zone_id = NULL WHERE id = %s",
|
||||
(MINER_UNASSIGNED,),
|
||||
)
|
||||
cur.execute(
|
||||
"UPDATE miners SET site_id = %s, agent_zone_id = %s WHERE id = %s",
|
||||
(SITE_B, ZONE_B, MINER_OTHER),
|
||||
)
|
||||
cur.execute(
|
||||
"UPDATE minerops_site_agents SET site_id = %s, agent_zone_id = %s "
|
||||
"WHERE agent_id = %s",
|
||||
(SITE_A, ZONE_A, AGENT_A),
|
||||
)
|
||||
connection.commit()
|
||||
finally:
|
||||
connection.close()
|
||||
|
||||
def test_repository_lists_multiple_canonical_networks(self):
|
||||
connection = connect()
|
||||
try:
|
||||
with connection.cursor() as cur:
|
||||
cur.execute(
|
||||
"INSERT INTO agent_zone_network_ranges (agent_zone_id, network) "
|
||||
"VALUES (%s, '10.30.0.0/24'), (%s, '10.31.0.0/24') "
|
||||
"ON CONFLICT DO NOTHING",
|
||||
(ZONE_A, ZONE_A),
|
||||
)
|
||||
connection.commit()
|
||||
finally:
|
||||
connection.close()
|
||||
listed = AgentZoneRepository(connect).list_zones(SITE_A)
|
||||
self.assertEqual([item["network"] for item in listed[0]["networks"]],
|
||||
["10.30.0.0/24", "10.31.0.0/24"])
|
||||
|
||||
def test_active_names_are_case_insensitively_unique_per_site(self):
|
||||
connection = connect()
|
||||
try:
|
||||
with self.assertRaises(UniqueViolation), connection.cursor() as cur:
|
||||
cur.execute(
|
||||
"INSERT INTO agent_zones (site_id, name) VALUES (%s, 'main room')",
|
||||
(SITE_A,),
|
||||
)
|
||||
finally:
|
||||
connection.rollback()
|
||||
connection.close()
|
||||
|
||||
def test_one_agent_per_zone_but_multiple_zones_per_site(self):
|
||||
connection = connect()
|
||||
try:
|
||||
with self.assertRaises(UniqueViolation), connection.cursor() as cur:
|
||||
cur.execute(
|
||||
"INSERT INTO minerops_site_agents "
|
||||
"(agent_id, site_id, agent_zone_id, site_name) "
|
||||
"VALUES ('agent-zone-postgres-duplicate', %s, %s, 'Duplicate')",
|
||||
(SITE_A, ZONE_A),
|
||||
)
|
||||
finally:
|
||||
connection.rollback()
|
||||
connection.close()
|
||||
|
||||
def test_composite_foreign_keys_reject_cross_site_assignment(self):
|
||||
connection = connect()
|
||||
try:
|
||||
with self.assertRaises(ForeignKeyViolation), connection.cursor() as cur:
|
||||
cur.execute(
|
||||
"UPDATE miners SET agent_zone_id = %s WHERE id = %s",
|
||||
(ZONE_A, MINER_OTHER),
|
||||
)
|
||||
finally:
|
||||
connection.rollback()
|
||||
connection.close()
|
||||
|
||||
def test_unassigned_miner_inherits_zone_and_existing_other_zone_wins(self):
|
||||
repository = AgentZoneRepository(connect)
|
||||
outcomes = repository.assign_miners_to_zone(
|
||||
[MINER_UNASSIGNED, MINER_OTHER], ZONE_A
|
||||
)
|
||||
self.assertEqual(outcomes[MINER_UNASSIGNED], "assigned")
|
||||
self.assertEqual(outcomes[MINER_OTHER], "agent_zone_conflict")
|
||||
connection = connect()
|
||||
try:
|
||||
with connection.cursor() as cur:
|
||||
cur.execute(
|
||||
"SELECT id, agent_zone_id FROM miners WHERE id IN (%s, %s)",
|
||||
(MINER_UNASSIGNED, MINER_OTHER),
|
||||
)
|
||||
assignments = dict(cur.fetchall())
|
||||
finally:
|
||||
connection.close()
|
||||
self.assertEqual(assignments[MINER_UNASSIGNED], ZONE_A)
|
||||
self.assertEqual(assignments[MINER_OTHER], ZONE_B)
|
||||
|
||||
def test_disabled_zone_blocks_control_without_site_fallback(self):
|
||||
connection = connect()
|
||||
try:
|
||||
with connection.cursor() as cur:
|
||||
cur.execute(
|
||||
"UPDATE miners SET agent_zone_id = %s WHERE id = %s",
|
||||
(ZONE_A, MINER_UNASSIGNED),
|
||||
)
|
||||
cur.execute("UPDATE agent_zones SET is_enabled = false WHERE id = %s", (ZONE_A,))
|
||||
connection.commit()
|
||||
finally:
|
||||
connection.close()
|
||||
try:
|
||||
projection = MinerControlService(MinerControlRepository(connect)).get_controls(
|
||||
MINER_UNASSIGNED
|
||||
)
|
||||
self.assertEqual(projection.service_control.reason, "agent_zone_disabled")
|
||||
finally:
|
||||
connection = connect()
|
||||
try:
|
||||
with connection.cursor() as cur:
|
||||
cur.execute("UPDATE agent_zones SET is_enabled = true WHERE id = %s", (ZONE_A,))
|
||||
connection.commit()
|
||||
finally:
|
||||
connection.close()
|
||||
|
||||
|
||||
def test_enabled_zone_without_agent_does_not_fall_back_within_site(self):
|
||||
connection = connect()
|
||||
try:
|
||||
with connection.cursor() as cur:
|
||||
cur.execute(
|
||||
"UPDATE miners SET agent_zone_id = %s WHERE id = %s",
|
||||
(ZONE_C, MINER_UNASSIGNED),
|
||||
)
|
||||
connection.commit()
|
||||
finally:
|
||||
connection.close()
|
||||
projection = MinerControlService(MinerControlRepository(connect)).get_controls(
|
||||
MINER_UNASSIGNED
|
||||
)
|
||||
self.assertEqual(projection.service_control.reason, "agent_unavailable")
|
||||
|
||||
def test_null_zone_does_not_fall_back_to_site_agent(self):
|
||||
projection = MinerControlService(MinerControlRepository(connect)).get_controls(
|
||||
MINER_UNASSIGNED
|
||||
)
|
||||
self.assertEqual(projection.service_control.reason, "agent_zone_unassigned")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in new issue
Block a user