AI_engine/tests/test_express_dispatch_agent.py

"""
Unit tests for the ExpressDispatchAgent — the DoormileExpress batch flow.

Focus, in the spirit of the other agent tests: the distribution must respect the
radius guard, the per-rider cap, and load balancing; and sequencing must survive
an optimizer that is down, wrong, or echoing ids we never sent, without ever
raising or writing a stop for a booking we did not assign.

The agent is built with __new__ so NATS / the message bus are never touched, and
the optimizer HTTP call is monkeypatched — no network.

Run:
    python -m unittest discover -s tests
"""
import asyncio
import unittest

import agents.express_dispatch_agent as eda
from agents.express_dispatch_agent import ExpressDispatchAgent


def _bare_agent():
    # Skip __init__ (Redis/NATS/registration); we only exercise pure decision code.
    return ExpressDispatchAgent.__new__(ExpressDispatchAgent)


def _booking(bid, plat=11.0045, plon=76.9612, dlat=11.05, dlon=76.93):
    return {
        "booking_id": bid, "booking_no": f"DM-{bid}",
        "pickuplatitude": plat, "pickuplongitude": plon,
        "deliverylatitude": dlat, "deliverylongitude": dlon,
    }


def _rider(mid, lat, lon):
    return {"miler_user_id": mid, "latitude": lat, "longitude": lon}


class DistributionTests(unittest.TestCase):
    def setUp(self):
        self.agent = _bare_agent()

    def test_splits_load_across_near_riders(self):
        bookings = [_booking(58), _booking(59), _booking(60)]
        riders = [_rider(38, 11.01, 76.96), _rider(41, 11.00, 76.97)]
        out = self.agent._distribute(bookings, riders)
        # Both riders used; nobody gets all three (load penalty spreads them).
        self.assertEqual(set(out.keys()), {38, 41})
        self.assertTrue(all(1 <= len(v) <= 2 for v in out.values()))
        self.assertEqual(sum(len(v) for v in out.values()), 3)

    def test_radius_guard_excludes_far_rider(self):
        bookings = [_booking(58), _booking(59)]
        # 38 is in Coimbatore; 99 is in Nagercoil (~200 km away) and must be excluded.
        riders = [_rider(38, 11.01, 76.96), _rider(99, 8.18, 77.43)]
        out = self.agent._distribute(bookings, riders)
        self.assertEqual(set(out.keys()), {38})
        self.assertEqual(99 in out, False)

    def test_no_gps_rider_assumed_at_pickup(self):
        bookings = [_booking(58)]
        riders = [_rider(38, 0, 0)]  # freshly Available, no fix yet
        out = self.agent._distribute(bookings, riders)
        self.assertIn(38, out)  # not excluded by the radius guard
        self.assertEqual(len(out[38]), 1)
        self.assertEqual(out[38][0]["booking_id"], 58)

    def test_per_rider_cap_leaves_overflow_unassigned(self):
        # 7 bookings, one rider, cap 5 -> 5 assigned, 2 left for manual assignment.
        old = eda.MAX_PER_RIDER
        eda.MAX_PER_RIDER = 5
        try:
            bookings = [_booking(i) for i in range(7)]
            riders = [_rider(38, 11.01, 76.96)]
            out = self.agent._distribute(bookings, riders)
            self.assertEqual(len(out[38]), 5)
        finally:
            eda.MAX_PER_RIDER = old


class SequenceTests(unittest.IsolatedAsyncioTestCase):
    def setUp(self):
        self.agent = _bare_agent()

    async def test_single_stop_is_step_one_without_calling_optimizer(self):
        called = []

        async def fake_post(*a, **k):
            called.append(1)
            return None

        eda.api_post = fake_post
        rows = await self.agent._sequence(13, _rider(38, 11.01, 76.96), [_booking(58)])
        self.assertEqual(called, [])  # no HTTP for a single stop
        self.assertEqual(len(rows), 1)
        self.assertEqual(rows[0]["step"], 1)
        self.assertEqual(rows[0]["booking_id"], 58)

    async def test_maps_optimizer_steps_back_by_bookingid(self):
        async def fake_post(url, **k):
            return {"success": True, "stops": [
                {"bookingid": 60, "step": 1, "previouskms": 0, "cumulativekms": 0, "etaminutes": 0, "cumulativeeta": 0},
                {"bookingid": 58, "step": 2, "previouskms": 4, "cumulativekms": 4, "etaminutes": 12, "cumulativeeta": 12},
                {"bookingid": 59, "step": 3, "previouskms": 5, "cumulativekms": 9, "etaminutes": 9, "cumulativeeta": 21},
            ]}
        eda.api_post = fake_post
        stops = [_booking(58), _booking(59), _booking(60)]
        rows = await self.agent._sequence(13, _rider(38, 11.01, 76.96), stops)
        by_id = {r["booking_id"]: r for r in rows}
        self.assertEqual(by_id[60]["step"], 1)
        self.assertEqual(by_id[58]["step"], 2)
        self.assertEqual(by_id[59]["cumulativeeta"], 21)

    async def test_optimizer_down_assigns_unsequenced_step_zero(self):
        async def fake_post(url, **k):
            return None  # 5xx / network -> http client returns None
        eda.api_post = fake_post
        stops = [_booking(58), _booking(59)]
        rows = await self.agent._sequence(13, _rider(38, 11.01, 76.96), stops)
        self.assertEqual(len(rows), 2)
        self.assertTrue(all(r["step"] == 0 for r in rows))  # assigned, just unordered

    async def test_unknown_bookingid_never_leaks_but_stop_still_assigned(self):
        async def fake_post(url, **k):
            return {"success": True, "stops": [
                {"bookingid": 999, "step": 1},   # never sent — must be dropped
                {"bookingid": 58, "step": 2},
            ]}
        eda.api_post = fake_post
        stops = [_booking(58), _booking(59)]
        rows = await self.agent._sequence(13, _rider(38, 11.01, 76.96), stops)
        ids = {r["booking_id"] for r in rows}
        self.assertNotIn(999, ids)          # echoed unknown id discarded
        self.assertEqual(ids, {58, 59})     # both real stops still assigned
        by_id = {r["booking_id"]: r for r in rows}
        self.assertEqual(by_id[58]["step"], 2)
        self.assertEqual(by_id[59]["step"], 0)  # dropped by optimizer -> unsequenced


if __name__ == "__main__":
    unittest.main()