diff --git a/robosystems_client/clients/operation_client.py b/robosystems_client/clients/operation_client.py index da05656..ef0704b 100644 --- a/robosystems_client/clients/operation_client.py +++ b/robosystems_client/clients/operation_client.py @@ -5,7 +5,7 @@ import logging from dataclasses import dataclass -from typing import Dict, Any, Optional, Callable, List +from typing import Dict, Any, Optional, Callable, List, cast from datetime import datetime from enum import Enum @@ -69,6 +69,24 @@ class MonitorOptions: poll_interval: Optional[int] = None +def _parsed_dict(parsed: Any) -> Dict[str, Any] | None: + """Body of a generated operations response as a plain dict. + + The operations endpoints are declared as free-form objects, so the + generator emits models whose only field is ``additional_properties`` + — attribute access on them raises, whatever the payload contained. + ``to_dict()`` is the accessor that works, and is what + ``operator_client._poll_for_completion`` already uses. + """ + if parsed is None: + return None + if hasattr(parsed, "to_dict"): + return cast(Dict[str, Any], parsed.to_dict()) + if isinstance(parsed, dict): + return cast(Dict[str, Any], parsed) + return None + + class OperationClient: """Client for monitoring operations via SSE""" @@ -218,23 +236,33 @@ def get_operation_status(self, operation_id: str) -> Dict[str, Any]: from ..api.operations.get_operation_status import ( sync_detailed as get_operation_status, ) - from ..client import Client + from .retry import retrying_client # Plain Client with the headers current now (`token_provider` wins). - client = Client(base_url=self.base_url, headers=resolve_auth_headers(self.config)) + client = retrying_client( + base_url=self.base_url, + headers=resolve_auth_headers(self.config), + config=self.config, + ) try: # Auth travels in self.headers (X-API-Key / Authorization). The generated # function takes no `token` kwarg — passing one raised TypeError, which # the handler below laundered into a fake {"status": "error"} result, so # this call never succeeded whenever a token was configured. + # + # Read through `to_dict()`, not attributes: the second incarnation of + # that same bug was `response.parsed.status` raising AttributeError on + # an additional-properties model, laundered into the same fake result + # for *every* response — including successful ones. response = get_operation_status(operation_id=operation_id, client=client) - if response.parsed: + payload = _parsed_dict(response.parsed) + if payload is not None: return { "operation_id": operation_id, - "status": response.parsed.status, - "progress": getattr(response.parsed, "progress", None), - "result": getattr(response.parsed, "result", None), - "error": getattr(response.parsed, "error", None), + "status": payload.get("status", "unknown"), + "progress": payload.get("progress"), + "result": payload.get("result"), + "error": payload.get("error"), } except Exception as e: # Logged rather than silently shaped into a result: swallowing here is @@ -248,26 +276,34 @@ def cancel_operation(self, operation_id: str) -> bool: """Cancel an operation""" # This would use the generated SDK to call /v1/operations/{operation_id}/cancel from ..api.operations.cancel_operation import sync_detailed as cancel_operation - from ..client import Client + from .retry import retrying_client # Plain Client with the headers current now (`token_provider` wins). - client = Client(base_url=self.base_url, headers=resolve_auth_headers(self.config)) + client = retrying_client( + base_url=self.base_url, + headers=resolve_auth_headers(self.config), + config=self.config, + ) try: - # See get_operation_status: no `token` kwarg on the generated function. + # See get_operation_status: no `token` kwarg, and the body is read + # through `to_dict()` because attribute access always raises. response = cancel_operation(operation_id=operation_id, client=client) - if response.parsed: - return response.parsed.cancelled or False + payload = _parsed_dict(response.parsed) + cancelled = bool(payload.get("cancelled")) if payload is not None else False except Exception as e: logger.warning("Failed to cancel operation %s: %s", operation_id, e) return False - # Also close any active SSE connection with thread safety + # Close any active SSE connection with thread safety. This used to sit + # after an early `return` on the success path, so the one case that + # needs it — the cancel actually landed — was the one case that skipped + # it, leaking the stream. with self._lock: if operation_id in self.active_operations: self.active_operations[operation_id].close() del self.active_operations[operation_id] - return False + return cancelled def list_operations(self) -> List[Dict[str, Any]]: """List all operations (if supported by the API)""" diff --git a/robosystems_client/clients/operator_client.py b/robosystems_client/clients/operator_client.py index 3bc35d5..f7e790b 100644 --- a/robosystems_client/clients/operator_client.py +++ b/robosystems_client/clients/operator_client.py @@ -16,6 +16,7 @@ sync_detailed as get_operation_status, ) from ..client import Client +from .retry import retrying_client from ..models.operator_request import OperatorRequest from ..models.operator_message import OperatorMessage from ..types import UNSET @@ -159,7 +160,11 @@ def _rest_client(self) -> Client: """ if not resolve_config_token(self.config): raise Exception("No API key provided. Set X-API-Key in headers.") - return Client(base_url=self.base_url, headers=resolve_auth_headers(self.config)) + return retrying_client( + base_url=self.base_url, + headers=resolve_auth_headers(self.config), + config=self.config, + ) def _sse_config(self) -> SSEConfig: """Stream config for one connect; headers carry the credential current now.""" diff --git a/robosystems_client/clients/query_client.py b/robosystems_client/clients/query_client.py index 6765a99..faa52e6 100644 --- a/robosystems_client/clients/query_client.py +++ b/robosystems_client/clients/query_client.py @@ -20,6 +20,7 @@ from ..api.query.execute_cypher import sync_detailed as execute_cypher_query from ..models.cypher_statement_request import CypherStatementRequest from ..client import Client +from .retry import retrying_client from .sse_client import ( SSEClient, AsyncSSEClient, @@ -97,7 +98,11 @@ def _rest_client(self) -> Client: """A REST client for one call, carrying the credential current now.""" if not resolve_config_token(self.config): raise Exception("No API key provided. Set X-API-Key in headers.") - return Client(base_url=self.base_url, headers=resolve_auth_headers(self.config)) + return retrying_client( + base_url=self.base_url, + headers=resolve_auth_headers(self.config), + config=self.config, + ) def _sse_config(self) -> SSEConfig: """Stream config for one connect; headers carry the credential current now.""" diff --git a/robosystems_client/clients/retry.py b/robosystems_client/clients/retry.py index 8ee51ee..99125ca 100644 --- a/robosystems_client/clients/retry.py +++ b/robosystems_client/clients/retry.py @@ -22,7 +22,7 @@ import httpx -from ..client import AuthenticatedClient +from ..client import AuthenticatedClient, Client RETRY_STATUS_CODES = frozenset({429}) @@ -174,3 +174,28 @@ def retrying_authenticated_client( base_url=base_url, headers=request_headers, timeout=None, config=config ) ) + + +def retrying_client( + *, + base_url: str, + headers: dict[str, str] | None = None, + config: dict[str, Any] | None = None, +) -> Client: + """A plain :class:`Client` whose transport replays 429s. + + The unauthenticated sibling of :func:`retrying_authenticated_client`, + for the facades that resolve their credential themselves and pass it + in ``headers`` (query / operator / operations) rather than letting + ``AuthenticatedClient`` stamp it. Those paths are rate-limited like + any other — Cypher queries, operator runs and operation-status polls + each draw on their own category budget — so they need the same replay. + """ + return Client( + base_url=base_url, + headers=dict(headers or {}), + ).set_httpx_client( + build_httpx_client( + base_url=base_url, headers=headers or {}, timeout=None, config=config + ) + ) diff --git a/tests/test_auth_header_resolution.py b/tests/test_auth_header_resolution.py index 2fa0abd..9547b62 100644 --- a/tests/test_auth_header_resolution.py +++ b/tests/test_auth_header_resolution.py @@ -13,6 +13,7 @@ from robosystems_client.clients.auth_integration import _apply_auth_header from robosystems_client.clients.operation_client import OperationClient +from robosystems_client.clients.retry import RetryingClient from robosystems_client.clients.sse_client import SSEClient, event_error_message from robosystems_client.clients.token_utils import ( apply_auth_header, @@ -137,9 +138,11 @@ def test_stream_headers_resolved_from_provider_at_connect( assert MockSSE.call_args[0][0].headers == {"Authorization": "Bearer jwt-rotated"} - # `Client` is imported inside the method, so patch it at its source module. - @patch("robosystems_client.client.Client") - def test_status_call_uses_provider_credential(self, MockClient, mock_config): + def test_status_call_uses_provider_credential(self, mock_config): + # Asserts the credential on the client actually handed to the + # generated op, rather than on a patched constructor — the status + # call now builds its client through `retrying_client`, and a + # constructor patch would only pin today's construction path. config = {**mock_config, "token_provider": lambda: "rfs_fresh"} with patch( "robosystems_client.api.operations.get_operation_status.sync_detailed" @@ -147,4 +150,6 @@ def test_status_call_uses_provider_credential(self, MockClient, mock_config): mock_get.return_value.parsed = None OperationClient(config).get_operation_status("op-1") - assert MockClient.call_args.kwargs["headers"] == {"X-API-Key": "rfs_fresh"} + passed = mock_get.call_args.kwargs["client"] + assert passed.get_httpx_client().headers["X-API-Key"] == "rfs_fresh" + assert isinstance(passed.get_httpx_client(), RetryingClient) diff --git a/tests/test_operation_client_ops.py b/tests/test_operation_client_ops.py index 823af89..045ba76 100644 --- a/tests/test_operation_client_ops.py +++ b/tests/test_operation_client_ops.py @@ -16,6 +16,12 @@ MonitorOptions, ) from robosystems_client.clients.sse_client import SSEClient +from robosystems_client.models.cancel_operation_response_canceloperation import ( + CancelOperationResponseCanceloperation, +) +from robosystems_client.models.get_operation_status_response_getoperationstatus import ( + GetOperationStatusResponseGetoperationstatus, +) # ── Helpers ────────────────────────────────────────────────────────── @@ -243,12 +249,14 @@ class TestGetOperationStatus: @patch("robosystems_client.api.operations.get_operation_status.sync_detailed") def test_get_status_success(self, mock_get, mock_config): """Test successful status retrieval.""" + # The real model, not a bare Mock: the endpoint is a free-form object, + # so the generated model carries only `additional_properties` and + # attribute access on it raises. A Mock makes `.status` work and let + # this test pass against code that could never work in production. mock_resp = Mock() - mock_resp.parsed = Mock() - mock_resp.parsed.status = "running" - mock_resp.parsed.progress = 50 - mock_resp.parsed.result = None - mock_resp.parsed.error = None + mock_resp.parsed = GetOperationStatusResponseGetoperationstatus.from_dict( + {"status": "running", "progress": 50, "result": None, "error": None} + ) mock_get.return_value = mock_resp client = OperationClient(mock_config) @@ -292,8 +300,9 @@ class TestCancelOperation: def test_cancel_success(self, mock_cancel, mock_config): """Test successful cancellation.""" mock_resp = Mock() - mock_resp.parsed = Mock() - mock_resp.parsed.cancelled = True + mock_resp.parsed = CancelOperationResponseCanceloperation.from_dict( + {"cancelled": True} + ) mock_cancel.return_value = mock_resp client = OperationClient(mock_config) diff --git a/tests/test_operation_client_status.py b/tests/test_operation_client_status.py new file mode 100644 index 0000000..aba3362 --- /dev/null +++ b/tests/test_operation_client_status.py @@ -0,0 +1,158 @@ +"""Unit tests for OperationClient's status and cancel calls. + +Both read the response body of a free-form operations endpoint. The +generator emits models whose only field is ``additional_properties``, so +attribute access on them raises whatever the payload held — and both +methods wrap the call in a broad ``except``, which turned that +AttributeError into a plausible-looking failure result for *every* +response, successful ones included. + +``operator_client._poll_for_completion`` already read these through +``to_dict()``; these tests pin that the same accessor is used here. +""" + +from typing import Any +from unittest.mock import MagicMock, patch + +import pytest + +from robosystems_client.clients.operation_client import ( + OperationClient, + _parsed_dict, +) +from robosystems_client.models.cancel_operation_response_canceloperation import ( + CancelOperationResponseCanceloperation, +) +from robosystems_client.models.get_operation_status_response_getoperationstatus import ( + GetOperationStatusResponseGetoperationstatus, +) + + +@pytest.fixture +def client() -> OperationClient: + return OperationClient( + {"base_url": "http://localhost:8000", "token": "rfs_test", "headers": {}} + ) + + +def _status_response(body: dict[str, Any]) -> MagicMock: + response = MagicMock() + response.parsed = GetOperationStatusResponseGetoperationstatus.from_dict(body) + return response + + +def _cancel_response(body: dict[str, Any]) -> MagicMock: + response = MagicMock() + response.parsed = CancelOperationResponseCanceloperation.from_dict(body) + return response + + +@pytest.mark.unit +class TestParsedDict: + def test_reads_an_additional_properties_model(self): + parsed = GetOperationStatusResponseGetoperationstatus.from_dict( + {"status": "completed"} + ) + # The accessor that motivated this helper: attribute access does not + # work on these models, whatever the payload contained. + assert not hasattr(parsed, "status") + assert _parsed_dict(parsed) == {"status": "completed"} + + def test_passes_a_plain_dict_through(self): + assert _parsed_dict({"status": "queued"}) == {"status": "queued"} + + def test_none_when_there_is_no_body(self): + assert _parsed_dict(None) is None + + +@pytest.mark.unit +class TestGetOperationStatus: + @patch("robosystems_client.api.operations.get_operation_status.sync_detailed") + def test_returns_the_real_status(self, mock_get, client: OperationClient): + mock_get.return_value = _status_response( + { + "operation_id": "op_1", + "status": "completed", + "progress": 100, + "result": {"rows": 3}, + } + ) + + result = client.get_operation_status("op_1") + + # Previously this whole payload was discarded and the method returned + # {"status": "error", "error": "...has no attribute 'status'"}. + assert result["status"] == "completed" + assert result["progress"] == 100 + assert result["result"] == {"rows": 3} + assert result["error"] is None + assert result["operation_id"] == "op_1" + + @patch("robosystems_client.api.operations.get_operation_status.sync_detailed") + def test_surfaces_a_failed_operation(self, mock_get, client: OperationClient): + mock_get.return_value = _status_response( + {"status": "failed", "error": "materialization timed out"} + ) + + result = client.get_operation_status("op_1") + + assert result["status"] == "failed" + assert result["error"] == "materialization timed out" + + @patch("robosystems_client.api.operations.get_operation_status.sync_detailed") + def test_unknown_when_the_body_omits_status(self, mock_get, client: OperationClient): + mock_get.return_value = _status_response({"operation_id": "op_1"}) + + assert client.get_operation_status("op_1")["status"] == "unknown" + + @patch("robosystems_client.api.operations.get_operation_status.sync_detailed") + def test_unknown_when_there_is_no_body(self, mock_get, client: OperationClient): + response = MagicMock() + response.parsed = None + mock_get.return_value = response + + assert client.get_operation_status("op_1")["status"] == "unknown" + + @patch("robosystems_client.api.operations.get_operation_status.sync_detailed") + def test_transport_failure_still_reports_an_error( + self, mock_get, client: OperationClient + ): + mock_get.side_effect = RuntimeError("connection refused") + + result = client.get_operation_status("op_1") + + assert result["status"] == "error" + assert "connection refused" in result["error"] + + +@pytest.mark.unit +class TestCancelOperation: + @patch("robosystems_client.api.operations.cancel_operation.sync_detailed") + def test_returns_true_when_the_cancel_lands(self, mock_cancel, client): + mock_cancel.return_value = _cancel_response({"cancelled": True}) + + assert client.cancel_operation("op_1") is True + + @patch("robosystems_client.api.operations.cancel_operation.sync_detailed") + def test_returns_false_when_the_server_declines(self, mock_cancel, client): + mock_cancel.return_value = _cancel_response({"cancelled": False}) + + assert client.cancel_operation("op_1") is False + + @patch("robosystems_client.api.operations.cancel_operation.sync_detailed") + def test_closes_the_stream_on_a_successful_cancel(self, mock_cancel, client): + # The cleanup used to sit after an early `return` on the success path, + # so the one case that needs it skipped it and leaked the stream. + mock_cancel.return_value = _cancel_response({"cancelled": True}) + stream = MagicMock() + client.active_operations["op_1"] = stream + + assert client.cancel_operation("op_1") is True + stream.close.assert_called_once() + assert "op_1" not in client.active_operations + + @patch("robosystems_client.api.operations.cancel_operation.sync_detailed") + def test_transport_failure_returns_false(self, mock_cancel, client): + mock_cancel.side_effect = RuntimeError("connection refused") + + assert client.cancel_operation("op_1") is False diff --git a/tests/test_rate_limit_retry.py b/tests/test_rate_limit_retry.py index 75c8786..93c4c48 100644 --- a/tests/test_rate_limit_retry.py +++ b/tests/test_rate_limit_retry.py @@ -16,11 +16,14 @@ import pytest from robosystems_client.clients.ledger_client import LedgerClient +from robosystems_client.clients.operation_client import OperationClient +from robosystems_client.clients.query_client import QueryClient from robosystems_client.clients.retry import ( RetryingClient, backoff_seconds, retry_after_seconds, retrying_authenticated_client, + retrying_client, ) @@ -211,3 +214,65 @@ def test_ledger_write_survives_a_rate_limit_burst(self, stub: _Stub): assert stub.calls == 3 assert result.id == "evt_1" assert all(p.endswith("/operations/create-event-block") for p in stub.paths) + + +@pytest.mark.unit +class TestUnauthenticatedFacadeWiring: + """The facades that resolve their own credential and pass it in headers. + + query / operator / operations build a plain ``Client`` rather than an + ``AuthenticatedClient``, so they were missed by the first pass. Their + endpoints carry their own category budgets — Cypher queries, operator + runs, operation-status polls — and need the same replay. + """ + + def test_plain_client_retries_and_keeps_its_headers(self, stub: _Stub): + stub.fail_first = 2 + client = retrying_client( + base_url=stub.base_url, + headers={"X-API-Key": "rfs_test", "X-Trace": "1"}, + config={"max_retries": 5, "retry_delay": 1}, + ) + http = client.get_httpx_client() + response = http.post("/x", json={}) + + assert response.status_code == 200 + assert stub.calls == 3 + assert http.headers["X-API-Key"] == "rfs_test" + assert http.headers["X-Trace"] == "1" + + def test_operation_status_poll_survives_a_rate_limit_burst(self, stub: _Stub): + # Polling a long-running operation in a loop is exactly the shape + # that exhausts a category budget. + stub.fail_first = 2 + stub.body = b'{"operation_id": "op_1", "status": "completed", "progress": 100}' + client = OperationClient( + { + "base_url": stub.base_url, + "token": "rfs_test", + "headers": {}, + "max_retries": 5, + "retry_delay": 1, + } + ) + + result = client.get_operation_status("op_1") + + assert stub.calls == 3 + assert all(p.endswith("/v1/operations/op_1/status") for p in stub.paths) + assert result["status"] == "completed" + + def test_query_client_builds_a_retrying_rest_client(self, stub: _Stub): + client = QueryClient( + { + "base_url": stub.base_url, + "token": "rfs_test", + "headers": {}, + "max_retries": 3, + "retry_delay": 1, + } + ) + http = client._rest_client().get_httpx_client() + + assert isinstance(http, RetryingClient) + assert http.headers["X-API-Key"] == "rfs_test"