| """CPU checks for replica failover and request errors before live fault injection.""" |
| import asyncio |
| import unittest |
| import tempfile |
| from pathlib import Path |
| from unittest.mock import patch |
| from urllib.parse import urlsplit |
|
|
| import httpx |
|
|
| import ha_proxy as proxy |
|
|
|
|
| class ProxyChecks(unittest.IsolatedAsyncioTestCase): |
| async def asyncSetUp(self): |
| proxy.inflight = 0 |
| proxy.healthy.update({url: True for url in proxy.BACKENDS}) |
| proxy.cooldown.update({url: 0 for url in proxy.BACKENDS}) |
| self.body = {"text": "Premise: same state\nHypothesis: first"} |
| self.order = proxy.ordered_backends(self.body) |
| self.calls = [] |
| self.status = {} |
|
|
| def upstream(request): |
| url = str(request.url).removesuffix("/classify") |
| self.calls.append(url) |
| result = self.status.get(url, 200) |
| if isinstance(result, Exception): |
| raise result |
| return httpx.Response(result, json={"embedding": [0, 1, 0]}) |
|
|
| proxy.client = httpx.AsyncClient(transport=httpx.MockTransport(upstream)) |
| self.api = httpx.AsyncClient(transport=httpx.ASGITransport(app=proxy.app), base_url="http://test") |
|
|
| async def asyncTearDown(self): |
| await self.api.aclose() |
| await proxy.client.aclose() |
|
|
| async def test_fallback_on_failure(self): |
| self.status[self.order[0]] = 500 |
| response = await self.api.post("/classify", json=self.body) |
| self.assertEqual(response.status_code, 200) |
| self.assertEqual(self.calls, self.order) |
| self.assertFalse(proxy.healthy[self.order[0]]) |
|
|
| async def test_connection_failure_falls_back(self): |
| self.status[self.order[0]] = httpx.ConnectError("worker disappeared") |
| response = await self.api.post("/classify", json=self.body) |
| self.assertEqual(response.status_code, 200) |
| self.assertEqual(self.calls, self.order) |
|
|
| async def test_bad_input_is_not_retried(self): |
| self.status[self.order[0]] = 400 |
| response = await self.api.post("/classify", json=self.body) |
| self.assertEqual(response.status_code, 400) |
| self.assertEqual(self.calls, self.order[:1]) |
|
|
| async def test_overload_tries_other_worker(self): |
| self.status[self.order[0]] = 429 |
| response = await self.api.post("/classify", json=self.body) |
| self.assertEqual(response.status_code, 200) |
| self.assertTrue(proxy.healthy[self.order[0]]) |
|
|
| async def test_no_workers_returns_503(self): |
| proxy.healthy.update({url: False for url in proxy.BACKENDS}) |
| response = await self.api.post("/classify", json=self.body) |
| self.assertEqual(response.status_code, 503) |
| self.assertEqual(self.calls, []) |
|
|
| async def test_inflight_hang_fails_over_after_health_loss(self): |
| await proxy.client.aclose() |
|
|
| async def upstream(request): |
| url = str(request.url).removesuffix("/classify") |
| if url == self.order[0]: |
| proxy.healthy[url] = False |
| await asyncio.sleep(10) |
| return httpx.Response(200, json={"embedding": [0, 1, 0]}) |
|
|
| proxy.client = httpx.AsyncClient(transport=httpx.MockTransport(upstream)) |
| response = await asyncio.wait_for(self.api.post("/classify", json=self.body), timeout=3) |
| self.assertEqual(response.status_code, 200) |
| self.assertEqual(response.headers["X-OpenJev-Backend"], self.order[1]) |
|
|
| async def test_bad_json_and_backpressure_release_counter(self): |
| response = await self.api.post("/classify", content="{") |
| self.assertEqual(response.status_code, 400) |
| self.assertEqual(proxy.inflight, 0) |
| proxy.inflight = proxy.MAX_INFLIGHT |
| response = await self.api.post("/classify", json=self.body) |
| self.assertEqual(response.status_code, 429) |
|
|
| def test_shared_premise_affinity(self): |
| other = {"text": "Premise: same state\nHypothesis: second"} |
| self.assertEqual(proxy.ordered_backends(other), self.order) |
|
|
| async def test_warming_backend_not_routed_until_ready(self): |
| with tempfile.TemporaryDirectory() as directory, patch.object(proxy, "READY_DIR", directory): |
| await proxy.check_backends() |
| self.assertFalse(any(proxy.healthy.values())) |
| self.assertEqual(self.calls, []) |
| port = urlsplit(proxy.BACKENDS[0]).port |
| (Path(directory) / f"{port}.ready").write_text("123\n") |
| await proxy.check_backends() |
| self.assertTrue(proxy.healthy[proxy.BACKENDS[0]]) |
| self.assertEqual(sum(proxy.healthy.values()), 1) |
|
|
|
|
| if __name__ == "__main__": |
| unittest.main() |
|
|
