AI_engine/tests/test_dispatch_agent.py

"""
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` is the set of miler ids whose miler_status:<id> says Available;
    every other geo member is treated as present-but-off-duty (the stale-index case)."""
    def __init__(self, geo=None, counter_start=0, raise_on=(), available=None, 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._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._geo:
                return '{"userid": 0, "status": "Break"}'
            return None
        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_available_miler_within_km"], 10)
        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_available_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_available_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_available_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_available_miler_within_km"], 10)   # coverage still gathered
        self.assertEqual(facts["failures_today"], 0)             # counter degraded to 0
        self.assertEqual(count, 0)


class TestLiveness(unittest.IsolatedAsyncioTestCase):
    """The prod incident: 26 riders in the geo index, none on duty. The old
    check said 'rider within 20 km' and called it a coverage gap."""

    async def test_stale_index_entries_are_not_available_riders(self):
        redis = FakeRedis(geo=["4", "7", "12"], available=[])     # all present, all off-duty
        facts, _ = await make_agent(redis)._gather_assignment_facts("coimbatore", 11.0, 76.9)
        self.assertEqual(facts["nearest_available_miler_within_km"], "none within 30km")
        self.assertEqual(facts["milers_in_geo_index_within_30km"], 3)   # but the index knows them

    async def test_first_available_candidate_wins(self):
        redis = FakeRedis(geo=["4", "7", "12"], available=["12"])
        found = await make_agent(redis)._find_zone(11.0, 76.9, radius_km=10)
        self.assertEqual(found["miler_id"], "12")
        self.assertEqual(found["stale_candidates_skipped"], 2)

    async def test_missing_status_key_means_unavailable(self):
        redis = FakeRedis(geo=["99"], available=[])
        redis._geo = []            # nothing in geo → get() returns None for status
        self.assertFalse(await make_agent(redis)._miler_is_available("99"))


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):
        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 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()