143 lines
5.5 KiB
Python
143 lines
5.5 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."""
|
|
def __init__(self, geo=None, counter_start=0, raise_on=()):
|
|
self._geo = list(geo) if geo is not None else []
|
|
self._counter = counter_start
|
|
self._raise_on = set(raise_on)
|
|
self.incr_calls = []
|
|
self.expire_calls = []
|
|
|
|
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["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(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.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 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()
|