diff --git a/tests/unit/test_arax.py b/tests/unit/test_arax.py new file mode 100644 index 0000000..e9dbd41 --- /dev/null +++ b/tests/unit/test_arax.py @@ -0,0 +1,218 @@ +"""Tests for ``workers.arax.worker``. + +Focused on what happens to the status code ARAX answers with: it used to be +logged and then dropped -- every failure became the same ``{"status": "error"}`` +blob that the workflow reported as a success -- so nothing downstream had a +status code to report. +""" + +import json +import logging + +import httpx +import pytest + +from workers.arax.worker import ( + BAD_GATEWAY, + GATEWAY_TIMEOUT, + ARAXServiceError, + arax, +) + +logger = logging.getLogger(__name__) + +QUERY = { + "message": { + "query_graph": { + "nodes": {"a": {"ids": ["MONDO:0005148"]}, "b": {}}, + "edges": {"e0": {"subject": "a", "object": "b"}}, + } + } +} + +PATHFINDER_QUERY = { + "message": { + "query_graph": { + "nodes": {"a": {"ids": ["MONDO:0005148"]}, "b": {"ids": ["CHEBI:15365"]}}, + "paths": {"p0": {"subject": "a", "object": "b"}}, + } + } +} + +ARAX_RESPONSE = { + "message": { + "query_graph": QUERY["message"]["query_graph"], + "knowledge_graph": { + "nodes": {"MONDO:0005148": {}}, + "edges": {"e0": {"subject": "a", "object": "b"}}, + }, + "results": [], + } +} + + +def _task(): + return [ + "task_id", + { + "query_id": "query_id", + "response_id": "response_id", + "workflow": json.dumps([{"id": "arax"}]), + "log_level": "20", + "otel": json.dumps({}), + "metadata": json.dumps({}), + }, + ] + + +def _patch_db(mocker, message=None): + """Stub the two db calls the worker makes, returning the save mock.""" + mocker.patch( + "workers.arax.worker.get_message", + new_callable=mocker.AsyncMock, + return_value=message if message is not None else dict(QUERY), + ) + return mocker.patch( + "workers.arax.worker.save_message", + new_callable=mocker.AsyncMock, + ) + + +def _patch_post(mocker, response=None, side_effect=None): + return mocker.patch( + "httpx.AsyncClient.post", + new_callable=mocker.AsyncMock, + return_value=response, + side_effect=side_effect, + ) + + +def _patch_span(mocker): + span = mocker.MagicMock() + mocker.patch("workers.arax.worker.get_current_span", return_value=span) + return span + + +def _http_response(status_code, json_body=None, text=None): + """A real httpx.Response, so is_success/.json()/.content behave as in prod.""" + request = httpx.Request("POST", "https://arax.example/query") + if json_body is not None: + return httpx.Response(status_code, json=json_body, request=request) + return httpx.Response(status_code, text=text or "", request=request) + + +@pytest.mark.asyncio +async def test_successful_query_saves_response_and_advances_workflow(mocker): + save = _patch_db(mocker) + span = _patch_span(mocker) + _patch_post(mocker, _http_response(200, json_body=ARAX_RESPONSE)) + task = _task() + + await arax(task, logger) + + span.set_attribute.assert_any_call("arax.status_code", 200) + saved = save.await_args.args[1] + assert saved["message"]["results"] == [] + # Provenance is still injected on the way through. + assert saved["message"]["knowledge_graph"]["edges"]["e0"]["sources"] == [ + { + "resource_id": "infores:shepherd-arax", + "resource_role": "aggregator_knowledge_source", + "source_record_urls": None, + "upstream_resource_ids": ["infores:arax"], + } + ] + assert json.loads(task[1]["workflow"]) == [{"id": "arax"}] + + +@pytest.mark.parametrize("status_code", [400, 404, 429, 500, 502]) +@pytest.mark.asyncio +async def test_error_status_is_propagated(mocker, status_code): + """ARAX's own status code reaches the exception, the span and the response.""" + save = _patch_db(mocker) + span = _patch_span(mocker) + _patch_post(mocker, _http_response(status_code, text="upstream said no")) + + with pytest.raises(ARAXServiceError) as excinfo: + await arax(_task(), logger) + + assert excinfo.value.status_code == status_code + assert f"HTTP {status_code}" in str(excinfo.value) + assert "upstream said no" in str(excinfo.value) + span.set_attribute.assert_any_call("arax.status_code", status_code) + + saved = save.await_args.args[1] + assert saved["status"] == "Error" + assert f"[HTTP {status_code}]" in saved["description"] + # Still a TRAPI response for the query that was asked. + assert saved["message"]["query_graph"] == QUERY["message"]["query_graph"] + assert saved["message"]["results"] == [] + + +@pytest.mark.asyncio +async def test_error_body_is_truncated(mocker): + _patch_db(mocker) + _patch_span(mocker) + _patch_post(mocker, _http_response(500, text="x" * 5000)) + + with pytest.raises(ARAXServiceError) as excinfo: + await arax(_task(), logger) + + message = str(excinfo.value) + assert "x" * 500 in message and "x" * 501 not in message + + +@pytest.mark.asyncio +async def test_timeout_reports_gateway_timeout(mocker): + save = _patch_db(mocker) + span = _patch_span(mocker) + _patch_post(mocker, side_effect=httpx.ReadTimeout("timed out")) + + with pytest.raises(ARAXServiceError) as excinfo: + await arax(_task(), logger) + + assert excinfo.value.status_code == GATEWAY_TIMEOUT + span.set_attribute.assert_any_call("arax.status_code", GATEWAY_TIMEOUT) + assert f"[HTTP {GATEWAY_TIMEOUT}]" in save.await_args.args[1]["description"] + + +@pytest.mark.asyncio +async def test_transport_error_reports_bad_gateway(mocker): + save = _patch_db(mocker) + span = _patch_span(mocker) + _patch_post(mocker, side_effect=httpx.ConnectError("")) + + with pytest.raises(ARAXServiceError) as excinfo: + await arax(_task(), logger) + + assert excinfo.value.status_code == BAD_GATEWAY + # httpx.ConnectError stringifies to nothing, so the class name carries it. + assert "ConnectError" in str(excinfo.value) + span.set_attribute.assert_any_call("arax.status_code", BAD_GATEWAY) + assert f"[HTTP {BAD_GATEWAY}]" in save.await_args.args[1]["description"] + + +@pytest.mark.asyncio +async def test_unparseable_success_body_keeps_the_status_code(mocker): + save = _patch_db(mocker) + _patch_span(mocker) + _patch_post(mocker, _http_response(200, text="not json")) + + with pytest.raises(ARAXServiceError) as excinfo: + await arax(_task(), logger) + + assert excinfo.value.status_code == 200 + assert "[HTTP 200]" in save.await_args.args[1]["description"] + + +@pytest.mark.asyncio +async def test_pathfinder_query_is_routed_without_calling_arax(mocker): + save = _patch_db(mocker, message=dict(PATHFINDER_QUERY)) + post = _patch_post(mocker, _http_response(200, json_body=ARAX_RESPONSE)) + task = _task() + + await arax(task, logger) + + assert not post.called + assert not save.called + assert json.loads(task[1]["workflow"]) == [{"id": "arax.pathfinder"}] diff --git a/workers/arax/worker.py b/workers/arax/worker.py index 34a9b46..e501116 100644 --- a/workers/arax/worker.py +++ b/workers/arax/worker.py @@ -5,6 +5,7 @@ import logging import uuid import httpx +from opentelemetry.trace import get_current_span from shepherd_utils.inject_shepherd_arax_provenance import ( add_shepherd_arax_to_edge_sources, ) @@ -20,8 +21,47 @@ GROUP = "consumer" CONSUMER = str(uuid.uuid4())[:8] TASK_LIMIT = 10 +ARAX_TIMEOUT = 270 tracer = setup_tracer(STREAM) LOGGER = get_worker_logger(STREAM) +# How much of a failing ARAX response body goes into the error. The body of a +# non-2xx can be an arbitrarily large HTML error page, and this string ends up +# in the query's logs, so keep only the head of it. +ERROR_BODY_BYTES = 500 +# Used when ARAX never answered at all, so there is no status code of its own to +# pass on: we are a gateway in front of it, and these are the codes that say so. +BAD_GATEWAY = 502 +GATEWAY_TIMEOUT = 504 + + +class ARAXServiceError(Exception): + """The ARAX service did not return a usable TRAPI response. + + Carries the upstream HTTP status code so it survives the whole way out: + onto the task span as ``arax.status_code``, into the query's logs via + ``run_task_lifecycle``, and so into what the caller gets back. Previously + the status code was only logged and the failure was swallowed into a + non-TRAPI ``{"status": "error"}`` blob that the workflow then reported as a + successful response, which left nothing downstream with a status code to + report (see the same fix in ``arax_pathfinder``). + """ + + def __init__(self, message: str, status_code: int): + super().__init__(message) + self.status_code = status_code + + +def body_head(response: httpx.Response) -> str: + """The first ``ERROR_BODY_BYTES`` of a response body, for an error message.""" + try: + body = response.content[:ERROR_BODY_BYTES] + except Exception: + # Body not readable (streamed/closed response) -- the status code is + # still worth reporting on its own. + return "" + if not body: + return "" + return f": {body.decode('utf-8', 'replace')}" def is_pathfinder_query(message): @@ -44,6 +84,86 @@ def is_pathfinder_query(message): return len(qpaths) == 1 +async def call_arax(message: dict, logger: logging.Logger) -> dict: + """POST the message to the ARAX service and return its TRAPI response. + + Raises ``ARAXServiceError`` -- carrying ARAX's own status code -- for + anything that isn't a parseable 2xx, so the code reaches the span, the + query's logs and the response instead of being logged and dropped. + """ + message["submitter"] = "Shepherd" + logger.info(f"Get the message from db {message}") + headers = {"Content-Type": "application/json"} + span = get_current_span() + try: + async with httpx.AsyncClient(timeout=ARAX_TIMEOUT) as client: + response = await client.post( + settings.arax_url, json=message, headers=headers + ) + except httpx.TimeoutException as e: + span.set_attribute("arax.status_code", GATEWAY_TIMEOUT) + raise ARAXServiceError( + f"ARAX service at {settings.arax_url} did not respond within " + f"{ARAX_TIMEOUT}s: {type(e).__name__}", + GATEWAY_TIMEOUT, + ) from e + except Exception as e: + # httpx reports connect failures, TLS errors and protocol errors as + # distinct classes, several of which stringify to an empty message -- + # hence the type name alongside the message. + span.set_attribute("arax.status_code", BAD_GATEWAY) + raise ARAXServiceError( + f"Error occurred calling ARAX service at {settings.arax_url}: " + f"{type(e).__name__}: {e}", + BAD_GATEWAY, + ) from e + + status_code = response.status_code + span.set_attribute("arax.status_code", status_code) + logger.info(f"Status Code from ARAX response: {status_code}") + if not response.is_success: + raise ARAXServiceError( + f"ARAX service at {settings.arax_url} returned HTTP " + f"{status_code}{body_head(response)}", + status_code, + ) + + try: + result = response.json() + except Exception as e: + # A 2xx whose body isn't TRAPI JSON is still a failed lookup, and + # ARAX's status code is the most useful thing we know about it. + raise ARAXServiceError( + f"ARAX service at {settings.arax_url} returned HTTP {status_code} " + f"with a body that could not be parsed as JSON: {e}", + status_code, + ) from e + + return add_shepherd_arax_to_edge_sources(result) + + +def error_response(message: dict, error: ARAXServiceError) -> dict: + """A TRAPI response reporting a failed ARAX call. + + ``status``/``description`` are TRAPI Response fields, so the status code + lands somewhere the caller already parses rather than only in the logs. The + query graph is carried over and the result containers are emptied, so what + comes back is still a valid TRAPI response for the query that was asked. + """ + query_graph = {} + if isinstance(message.get("message"), dict): + query_graph = message["message"].get("query_graph") or {} + return { + "message": { + "query_graph": query_graph, + "knowledge_graph": {"nodes": {}, "edges": {}}, + "results": [], + }, + "status": "Error", + "description": f"[HTTP {error.status_code}] {error}", + } + + async def arax(task, logger: logging.Logger): query_id = task[1]["query_id"] logger.info(f"Getting message from db for query id {query_id}") @@ -51,21 +171,18 @@ async def arax(task, logger: logging.Logger): if is_pathfinder_query(message): task[1]["workflow"] = json.dumps([{"id": "arax.pathfinder"}]) else: - try: - message["submitter"] = "Shepherd" - logger.info(f"Get the message from db {message}") - headers = {"Content-Type": "application/json"} - async with httpx.AsyncClient(timeout=270) as client: - response = await client.post( - settings.arax_url, json=message, headers=headers - ) - logger.info(f"Status Code from ARAX response: {response.status_code}") - result = response.json() - result = add_shepherd_arax_to_edge_sources(result) - except Exception as e: - logger.error(f"Error occurred calling ARAX service: {e}") - result = {"status": "error", "error": str(e)} response_id = task[1]["response_id"] + try: + result = await call_arax(message, logger) + except ARAXServiceError as e: + # Leave the caller a TRAPI response that says what happened before + # letting the failure reach run_task_lifecycle, which records it on + # the span and routes the query to finish_query with an ERROR + # status. Without this the response id still holds the echo of the + # incoming query, so the caller gets a query that looks like it + # simply found nothing. + await save_message(response_id, error_response(message, e), logger) + raise await save_message(response_id, result, logger) task[1]["workflow"] = json.dumps([{"id": "arax"}])