|
20 | 20 |
|
21 | 21 | import os |
22 | 22 | import sys |
| 23 | +import asyncio |
| 24 | +import types |
23 | 25 | import unittest |
24 | 26 | from unittest import mock |
25 | 27 |
|
@@ -59,7 +61,24 @@ class TestRouterAdmissionLifecycle(unittest.TestCase): |
59 | 61 | """ |
60 | 62 |
|
61 | 63 | def _build_router(self): |
62 | | - from miles.router.router import MilesRouter |
| 64 | + ray_stub = types.ModuleType("ray") |
| 65 | + ray_stub.remote = lambda *args, **kwargs: ( |
| 66 | + args[0] if args and callable(args[0]) and not kwargs else lambda obj: obj |
| 67 | + ) |
| 68 | + ray_util_stub = types.ModuleType("ray.util") |
| 69 | + scheduling_stub = types.ModuleType("ray.util.scheduling_strategies") |
| 70 | + scheduling_stub.NodeAffinitySchedulingStrategy = object |
| 71 | + |
| 72 | + with mock.patch.dict( |
| 73 | + sys.modules, |
| 74 | + { |
| 75 | + "ray": ray_stub, |
| 76 | + "ray.util": ray_util_stub, |
| 77 | + "ray.util.scheduling_strategies": scheduling_stub, |
| 78 | + }, |
| 79 | + ): |
| 80 | + from miles.router.router import MilesRouter |
| 81 | + self.router_module = sys.modules[MilesRouter.__module__] |
63 | 82 |
|
64 | 83 | args = mock.Mock() |
65 | 84 | args.miles_router_max_connections = 8 |
@@ -96,6 +115,26 @@ def test_remove_worker_drops_all_state(self): |
96 | 115 | self.assertNotIn("http://w1:8000", router.enabled_workers) |
97 | 116 | self.assertNotIn("http://w1:8000", router.worker_engine_index_map) |
98 | 117 |
|
| 118 | + def test_health_check_does_not_probe_disabled_workers_when_zero_active(self): |
| 119 | + router = self._build_router() |
| 120 | + router._add_worker_internal("http://w1:8000", engine_index=0) |
| 121 | + router._disable_worker_internal("http://w1:8000") |
| 122 | + router._check_worker_health = mock.AsyncMock() |
| 123 | + |
| 124 | + sleep_calls = 0 |
| 125 | + |
| 126 | + async def sleep_once_then_cancel(_interval): |
| 127 | + nonlocal sleep_calls |
| 128 | + sleep_calls += 1 |
| 129 | + if sleep_calls > 1: |
| 130 | + raise asyncio.CancelledError |
| 131 | + |
| 132 | + with mock.patch.object(self.router_module.asyncio, "sleep", sleep_once_then_cancel): |
| 133 | + with self.assertRaises(asyncio.CancelledError): |
| 134 | + asyncio.run(router._health_check_loop()) |
| 135 | + |
| 136 | + router._check_worker_health.assert_not_called() |
| 137 | + |
99 | 138 |
|
100 | 139 | class TestSchedulerPreemptClassification(unittest.TestCase): |
101 | 140 | """F3 / F31 — _is_scheduler_preempt strict missing-metadata check.""" |
|
0 commit comments