333 lines
15 KiB
Python
333 lines
15 KiB
Python
"""
|
|
Unit tests for the DispatchAgent assignment-failure context gatherer.
|
|
|
|
Focus is the same as the ExceptionAgent tests: the gatherer must degrade
|
|
gracefully (no coords, a failing GEORADIUS, a failing counter) and never raise.
|
|
The real `_find_zone` is exercised through an in-memory fake Redis; the agent is
|
|
built with __new__ so NATS / the message bus are never touched.
|
|
|
|
Run:
|
|
python -m unittest discover -s tests
|
|
"""
|
|
import unittest
|
|
|
|
import nats.js.errors
|
|
|
|
from agents.dispatch_agent import DispatchAgent
|
|
|
|
|
|
class FakeRedis:
|
|
"""decode_responses=True style. `geo` is the GEORADIUS result for every call.
|
|
`available` / `on_break` are miler ids with a miler_status:<id> key saying so.
|
|
Any geo member in neither set has NO status key at all — the common case in
|
|
production, which must read as unknown presence, not off-duty."""
|
|
def __init__(self, geo=None, counter_start=0, raise_on=(), available=None,
|
|
on_break=(), alerted=False):
|
|
self._geo = list(geo) if geo is not None else []
|
|
self._counter = counter_start
|
|
self._raise_on = set(raise_on)
|
|
self._available = set(available) if available is not None else set(self._geo)
|
|
self._on_break = set(on_break)
|
|
self._kv = {}
|
|
if alerted:
|
|
self._kv["__alerted__"] = "1"
|
|
self.incr_calls = []
|
|
self.expire_calls = []
|
|
self.set_calls = []
|
|
|
|
async def get(self, key):
|
|
if key.startswith("miler_status:"):
|
|
mid = key.split(":", 1)[1]
|
|
if mid in self._available:
|
|
return '{"userid": %s, "status": "Available"}' % (mid if mid.isdigit() else 0)
|
|
if mid in self._on_break:
|
|
return '{"userid": 0, "status": "Break"}'
|
|
return None # no status key — unknown, not off-duty
|
|
return self._kv.get(key)
|
|
|
|
async def exists(self, key):
|
|
if key.startswith("assignment_alert_sent:"):
|
|
return 1 if "__alerted__" in self._kv or key in self._kv else 0
|
|
return 1 if key in self._kv else 0
|
|
|
|
async def set(self, key, value, ex=None, nx=False):
|
|
self.set_calls.append((key, ex, nx))
|
|
if nx and key in self._kv:
|
|
return None
|
|
self._kv[key] = value
|
|
return True
|
|
|
|
async def georadius(self, key, lon, lat, radius, unit, sort=None, count=None):
|
|
if "georadius" in self._raise_on:
|
|
raise RuntimeError("geo down")
|
|
return list(self._geo)
|
|
|
|
async def hgetall(self, key):
|
|
return {"hub_id": "H1", "zone_id": "z1", "avg_delivery_time": "60"}
|
|
|
|
async def incr(self, key):
|
|
if "incr" in self._raise_on:
|
|
raise RuntimeError("redis down")
|
|
self._counter += 1
|
|
self.incr_calls.append(key)
|
|
return self._counter
|
|
|
|
async def expire(self, key, ttl):
|
|
self.expire_calls.append((key, ttl))
|
|
return True
|
|
|
|
|
|
def make_agent(redis):
|
|
agent = DispatchAgent.__new__(DispatchAgent) # skip heavy __init__ / message-bus registration
|
|
agent._redis = redis
|
|
return agent
|
|
|
|
|
|
class FakeJS:
|
|
"""Stand-in JetStream context for binding tests."""
|
|
def __init__(self, stream_name=None, not_found=False):
|
|
self._stream = stream_name
|
|
self._not_found = not_found
|
|
self.subscribe_calls = []
|
|
self.add_stream_calls = []
|
|
|
|
async def find_stream_name_by_subject(self, subject):
|
|
if self._not_found:
|
|
raise nats.js.errors.NotFoundError()
|
|
return self._stream
|
|
|
|
async def subscribe(self, subject, durable=None, stream=None, cb=None):
|
|
self.subscribe_calls.append((subject, durable, stream))
|
|
return object() # stand-in subscription
|
|
|
|
async def add_stream(self, **kwargs):
|
|
self.add_stream_calls.append(kwargs)
|
|
|
|
|
|
async def _cb(msg):
|
|
pass
|
|
|
|
|
|
def make_agent_js(js):
|
|
agent = DispatchAgent.__new__(DispatchAgent)
|
|
agent._nats_js = js
|
|
return agent
|
|
|
|
|
|
class TestGatherAssignmentFacts(unittest.IsolatedAsyncioTestCase):
|
|
async def test_rider_nearby(self):
|
|
redis = FakeRedis(geo=["m5"], counter_start=0) # found on the first (10km) sweep
|
|
facts, count = await make_agent(redis)._gather_assignment_facts("hyderabad", 17.4, 78.4)
|
|
self.assertEqual(facts["zone_id"], "hyderabad")
|
|
self.assertTrue(facts["has_coordinates"])
|
|
self.assertEqual(facts["nearest_miler_within_km"], 10)
|
|
self.assertEqual(facts["nearest_miler_presence"], "available")
|
|
self.assertEqual(facts["milers_in_geo_index_within_30km"], 1)
|
|
self.assertFalse(facts["alert_already_sent_today"])
|
|
self.assertEqual(facts["failures_today"], 1)
|
|
self.assertEqual(count, 1)
|
|
self.assertEqual(redis.expire_calls[0][1], 172800)
|
|
|
|
async def test_no_rider_within_30km(self):
|
|
redis = FakeRedis(geo=[], counter_start=2) # every sweep empty
|
|
facts, count = await make_agent(redis)._gather_assignment_facts("pune", 18.5, 73.8)
|
|
self.assertEqual(facts["nearest_miler_within_km"], "none within 30km")
|
|
self.assertEqual(facts["milers_in_geo_index_within_30km"], 0)
|
|
self.assertEqual(count, 3)
|
|
|
|
async def test_no_coordinates(self):
|
|
redis = FakeRedis(geo=["m1"])
|
|
facts, count = await make_agent(redis)._gather_assignment_facts("unknown", None, None)
|
|
self.assertFalse(facts["has_coordinates"])
|
|
self.assertEqual(facts["nearest_miler_within_km"], "unknown (no coordinates)")
|
|
self.assertNotIn("milers_in_geo_index_within_30km", facts)
|
|
self.assertEqual(count, 1) # counter still incremented
|
|
self.assertEqual(redis.incr_calls and 1, 1)
|
|
|
|
async def test_georadius_error_is_swallowed(self):
|
|
redis = FakeRedis(raise_on={"georadius"})
|
|
facts, count = await make_agent(redis)._gather_assignment_facts("z", 1.0, 2.0) # must not raise
|
|
self.assertEqual(facts["nearest_miler_within_km"], "none within 30km")
|
|
self.assertEqual(facts["failures_today"], 1)
|
|
|
|
async def test_counter_error_is_swallowed(self):
|
|
redis = FakeRedis(geo=["m9"], raise_on={"incr"})
|
|
facts, count = await make_agent(redis)._gather_assignment_facts("z", 1.0, 2.0) # must not raise
|
|
self.assertEqual(facts["nearest_miler_within_km"], 10) # coverage still gathered
|
|
self.assertEqual(facts["failures_today"], 0) # counter degraded to 0
|
|
self.assertEqual(count, 0)
|
|
|
|
|
|
class TestPresence(unittest.IsolatedAsyncioTestCase):
|
|
"""Presence is three-state. The prod incident had 26 riders in the geo index;
|
|
only 2 of 34 milers had a status key at all, so "no key" must read as unknown
|
|
— reporting it as off-duty would be a confident wrong answer."""
|
|
|
|
async def test_confirmed_available_rider_is_preferred(self):
|
|
redis = FakeRedis(geo=["4", "7", "12"], available=["12"], on_break=["4", "7"])
|
|
found = await make_agent(redis)._find_zone(11.0, 76.9, radius_km=10)
|
|
self.assertEqual(found["miler_id"], "12")
|
|
self.assertEqual(found["presence"], "available")
|
|
|
|
async def test_no_status_key_reads_as_unknown_not_unavailable(self):
|
|
redis = FakeRedis(geo=["4", "7"], available=[], on_break=[]) # no status keys at all
|
|
agent = make_agent(redis)
|
|
self.assertEqual(await agent._miler_presence("4"), "unknown")
|
|
found = await agent._find_zone(11.0, 76.9, radius_km=10)
|
|
self.assertIsNotNone(found) # still reported — may be assignable
|
|
self.assertEqual(found["presence"], "unknown")
|
|
|
|
async def test_all_confirmed_off_duty_returns_none(self):
|
|
redis = FakeRedis(geo=["4", "7"], available=[], on_break=["4", "7"])
|
|
self.assertIsNone(await make_agent(redis)._find_zone(11.0, 76.9, radius_km=10))
|
|
|
|
async def test_facts_report_presence_unknown_with_rider_nearby(self):
|
|
redis = FakeRedis(geo=["4", "7", "12"], available=[], on_break=[])
|
|
facts, _ = await make_agent(redis)._gather_assignment_facts("coimbatore", 11.0, 76.9)
|
|
self.assertEqual(facts["nearest_miler_within_km"], 10)
|
|
self.assertEqual(facts["nearest_miler_presence"], "unknown")
|
|
self.assertEqual(facts["milers_in_geo_index_within_30km"], 3)
|
|
|
|
async def test_facts_report_none_found_when_all_off_duty(self):
|
|
redis = FakeRedis(geo=["4"], available=[], on_break=["4"])
|
|
facts, _ = await make_agent(redis)._gather_assignment_facts("coimbatore", 11.0, 76.9)
|
|
self.assertEqual(facts["nearest_miler_within_km"], "none within 30km")
|
|
self.assertEqual(facts["nearest_miler_presence"], "none found")
|
|
self.assertEqual(facts["milers_in_geo_index_within_30km"], 1) # index still knows them
|
|
|
|
|
|
class FakeMessages:
|
|
def __init__(self):
|
|
self.sent = []
|
|
|
|
async def send_message(self, recipient, message_type, payload, correlation_id=None):
|
|
self.sent.append((recipient, message_type, payload))
|
|
|
|
|
|
def make_handler_agent(redis, decision):
|
|
"""Agent with send_message captured and the LLM stubbed to a fixed decision (or None)."""
|
|
from agents import dispatch_agent as mod
|
|
agent = DispatchAgent.__new__(DispatchAgent)
|
|
agent._redis = redis
|
|
agent.agent_id = "DISPATCH_AGENT"
|
|
box = FakeMessages()
|
|
agent.send_message = box.send_message
|
|
calls = []
|
|
async def fake_decide(context, model=None):
|
|
calls.append(context)
|
|
return decision
|
|
mod.decide_assignment_failure = fake_decide
|
|
return agent, box, calls
|
|
|
|
|
|
class FakeMsg:
|
|
def __init__(self, payload):
|
|
import json
|
|
self.data = json.dumps(payload).encode()
|
|
self.acked = False
|
|
|
|
async def ack(self):
|
|
self.acked = True
|
|
|
|
|
|
class TestRateLimitAndSinks(unittest.IsolatedAsyncioTestCase):
|
|
def setUp(self):
|
|
from agents import dispatch_agent as mod
|
|
self._orig = mod.decide_assignment_failure
|
|
|
|
def tearDown(self):
|
|
from agents import dispatch_agent as mod
|
|
mod.decide_assignment_failure = self._orig
|
|
|
|
async def test_ops_alert_goes_to_jarvis_as_exception_and_marks_zone(self):
|
|
from core.types import MessageType
|
|
from core.llm import AssignmentDecision
|
|
redis = FakeRedis(geo=[], counter_start=2)
|
|
agent, box, calls = make_handler_agent(redis, AssignmentDecision("ops_alert", "gap", 0.9))
|
|
msg = FakeMsg({"booking_id": 1, "lat": 11.0, "lon": 76.9})
|
|
await agent._on_nats_booking_assignment_failed(msg)
|
|
self.assertTrue(msg.acked)
|
|
self.assertEqual(len(calls), 1) # LLM consulted once
|
|
recipient, mtype, payload = box.sent[0]
|
|
self.assertEqual(recipient, "JARVIS") # not CUSTOMER_AGENT
|
|
self.assertEqual(mtype, MessageType.EXCEPTION_DETECTED) # the path JARVIS handles
|
|
self.assertEqual(payload["exception_type"], "coverage_gap")
|
|
self.assertTrue(any(k.startswith("assignment_alert_sent:") for k, _, _ in redis.set_calls))
|
|
|
|
async def test_burst_after_alert_skips_llm_and_realerts_every_n(self):
|
|
from agents import dispatch_agent as mod
|
|
from core.llm import AssignmentDecision
|
|
redis = FakeRedis(geo=[], counter_start=0, alerted=True) # zone already alerted today
|
|
agent, box, calls = make_handler_agent(redis, AssignmentDecision("ops_alert", "gap", 0.9))
|
|
old = mod.DISPATCH_REALERT_EVERY
|
|
mod.DISPATCH_REALERT_EVERY = 100
|
|
try:
|
|
for i in range(1, 251):
|
|
await agent._on_nats_booking_assignment_failed(FakeMsg({"booking_id": i, "lat": 11.0, "lon": 76.9}))
|
|
finally:
|
|
mod.DISPATCH_REALERT_EVERY = old
|
|
self.assertEqual(calls, []) # zero LLM calls during the burst
|
|
self.assertEqual(len(box.sent), 2) # re-alerts at 100 and 200 only
|
|
self.assertIn("200 failed assignments", box.sent[1][2]["reasoning"])
|
|
|
|
async def test_llm_unavailable_falls_back_to_heuristic(self):
|
|
redis = FakeRedis(geo=[], counter_start=2)
|
|
agent, box, calls = make_handler_agent(redis, None) # LLM down
|
|
await agent._on_nats_booking_assignment_failed(FakeMsg({"booking_id": 5, "lat": 1.0, "lon": 2.0}))
|
|
self.assertEqual(len(box.sent), 1) # count reached 3 → heuristic alert
|
|
self.assertEqual(box.sent[0][0], "JARVIS")
|
|
|
|
async def test_escalate_goes_to_jarvis_as_exception(self):
|
|
from core.types import MessageType
|
|
from core.llm import AssignmentDecision
|
|
redis = FakeRedis(geo=["4"], available=["4"], counter_start=4)
|
|
agent, box, calls = make_handler_agent(redis, AssignmentDecision("escalate", "riders near, still failing", 0.8))
|
|
await agent._on_nats_booking_assignment_failed(FakeMsg({"booking_id": 9, "lat": 1.0, "lon": 2.0}))
|
|
recipient, mtype, payload = box.sent[0]
|
|
self.assertEqual((recipient, mtype), ("JARVIS", MessageType.EXCEPTION_DETECTED))
|
|
self.assertEqual(payload["proposed_action"], "escalate")
|
|
|
|
|
|
class TestZoneKey(unittest.IsolatedAsyncioTestCase):
|
|
"""The backend sends no zone_id today; without a fallback every failure in
|
|
the country shares one bucket and the daily alert limit becomes global."""
|
|
|
|
async def test_real_zone_id_wins(self):
|
|
from agents.dispatch_agent import zone_key
|
|
self.assertEqual(zone_key("hyderabad", 17.4, 78.4), "hyderabad")
|
|
|
|
async def test_missing_zone_falls_back_to_geo_cell(self):
|
|
from agents.dispatch_agent import zone_key
|
|
self.assertEqual(zone_key(None, 11.0182714, 76.9677744), "geo:11.0,77.0")
|
|
self.assertEqual(zone_key("unknown", 11.0182714, 76.9677744), "geo:11.0,77.0")
|
|
|
|
async def test_distant_pickups_get_different_buckets(self):
|
|
from agents.dispatch_agent import zone_key
|
|
self.assertNotEqual(zone_key(None, 11.01, 76.96), zone_key(None, 17.44, 78.39))
|
|
|
|
async def test_no_zone_and_no_coords_is_unknown(self):
|
|
from agents.dispatch_agent import zone_key
|
|
self.assertEqual(zone_key(None, None, None), "unknown")
|
|
self.assertEqual(zone_key("unknown", "bad", "coords"), "unknown")
|
|
|
|
|
|
class TestBindConsumer(unittest.IsolatedAsyncioTestCase):
|
|
async def test_binds_to_discovered_stream_without_creating(self):
|
|
js = FakeJS(stream_name="TRACKING")
|
|
sub = await make_agent_js(js)._bind_consumer("booking.assigned", "d1", _cb)
|
|
self.assertIsNotNone(sub)
|
|
# Bound explicitly on the discovered stream...
|
|
self.assertEqual(js.subscribe_calls, [("booking.assigned", "d1", "TRACKING")])
|
|
# ...and never tried to create a stream (the old overlap bug).
|
|
self.assertEqual(js.add_stream_calls, [])
|
|
|
|
async def test_returns_none_when_no_stream_carries_subject(self):
|
|
js = FakeJS(not_found=True)
|
|
sub = await make_agent_js(js)._bind_consumer("booking.assigned", "d1", _cb)
|
|
self.assertIsNone(sub) # clear failure, not a raise
|
|
self.assertEqual(js.subscribe_calls, []) # did not attempt to subscribe
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|