diff --git a/pyrit/cli/api_client.py b/pyrit/cli/api_client.py index ac24fdf454..47d587aa4a 100644 --- a/pyrit/cli/api_client.py +++ b/pyrit/cli/api_client.py @@ -92,7 +92,10 @@ async def health_check_async(self) -> bool: try: client = self._get_client() resp = await client.get("/api/health") - return resp.status_code == 200 + if resp.status_code != 200: + return False + payload = resp.json() + return payload.get("status") == "healthy" and payload.get("service") == "pyrit-backend" except httpx.ConnectError: return False except Exception: diff --git a/tests/unit/cli/test_api_client.py b/tests/unit/cli/test_api_client.py index 7c6341b1cb..f1b83a6c96 100644 --- a/tests/unit/cli/test_api_client.py +++ b/tests/unit/cli/test_api_client.py @@ -168,11 +168,22 @@ def test_get_client_raises_when_not_opened(): async def test_health_check_returns_true_on_200(client, mock_httpx_client): - mock_httpx_client.get.return_value = _make_response(status_code=200) + mock_httpx_client.get.return_value = _make_response( + status_code=200, + json_data={"status": "healthy", "service": "pyrit-backend"}, + ) assert await client.health_check_async() is True mock_httpx_client.get.assert_awaited_once_with("/api/health") +async def test_health_check_returns_false_for_unrelated_service(client, mock_httpx_client): + mock_httpx_client.get.return_value = _make_response( + status_code=200, + json_data={"status": "ok", "service": "another-service"}, + ) + assert await client.health_check_async() is False + + async def test_health_check_returns_false_on_non_200(client, mock_httpx_client): mock_httpx_client.get.return_value = _make_response(status_code=503) assert await client.health_check_async() is False