diff --git a/cueapi/__init__.py b/cueapi/__init__.py index 3dc5817..86ca880 100644 --- a/cueapi/__init__.py +++ b/cueapi/__init__.py @@ -12,6 +12,8 @@ ) from cueapi.payload import CuePayload from cueapi.resources.executions import ExecutionsResource +from cueapi.resources.usage import UsageResource +from cueapi.resources.workers import WorkersResource from cueapi.webhook import verify_webhook __version__ = "0.1.2" @@ -20,6 +22,8 @@ "CueAPI", "CuePayload", "ExecutionsResource", + "UsageResource", + "WorkersResource", "verify_webhook", "CueAPIError", "AuthenticationError", diff --git a/cueapi/client.py b/cueapi/client.py index 696d0fd..0fe27b4 100644 --- a/cueapi/client.py +++ b/cueapi/client.py @@ -17,6 +17,8 @@ ) from cueapi.resources.cues import CuesResource from cueapi.resources.executions import ExecutionsResource +from cueapi.resources.usage import UsageResource +from cueapi.resources.workers import WorkersResource DEFAULT_BASE_URL = "https://api.cueapi.ai" DEFAULT_TIMEOUT = 30.0 @@ -69,6 +71,8 @@ def __init__( # Resources self.cues = CuesResource(self) self.executions = ExecutionsResource(self) + self.workers = WorkersResource(self) + self.usage = UsageResource(self) def close(self) -> None: """Close the underlying HTTP client.""" diff --git a/cueapi/resources/usage.py b/cueapi/resources/usage.py new file mode 100644 index 0000000..230db14 --- /dev/null +++ b/cueapi/resources/usage.py @@ -0,0 +1,31 @@ +"""Usage resource — plan + cue + execution usage stats.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from cueapi.client import CueAPI + + +class UsageResource: + """Usage stats resource. + + Wraps ``GET /v1/usage`` for SDK callers who want plan + cue count + + execution count + rate-limit info without parsing the broader + ``/v1/auth/me`` response. + """ + + def __init__(self, client: "CueAPI") -> None: + self._client = client + + def get(self) -> dict: + """Get current usage stats. + + Returns: + Dict with ``plan`` (name + interval + period_end), + ``cues`` (active count + limit), + ``executions`` (used this period + limit + outcomes summary), + and ``rate_limit`` (requests/min limit). + """ + return self._client._get("/v1/usage") diff --git a/cueapi/resources/workers.py b/cueapi/resources/workers.py new file mode 100644 index 0000000..23e70f0 --- /dev/null +++ b/cueapi/resources/workers.py @@ -0,0 +1,51 @@ +"""Workers resource — fleet visibility for worker-transport users.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from cueapi.client import CueAPI + + +class WorkersResource: + """Workers API resource. + + Mirrors the hosted ``/v1/workers`` surface — list registered workers + with heartbeat status, and delete decommissioned workers. Worker + registration itself happens via cueapi-worker (which sends heartbeats); + the SDK doesn't expose ``POST /v1/worker/heartbeat`` because direct + SDK-driven registration is redundant with that package. + """ + + def __init__(self, client: "CueAPI") -> None: + self._client = client + + def list(self) -> dict: + """List all registered workers with heartbeat status. + + Returns: + Dict with ``workers`` (list of worker dicts) and ``total``. + Each worker carries ``worker_id``, ``handlers``, + ``last_heartbeat``, ``heartbeat_status`` + (``online`` / ``stale`` / ``dead`` based on seconds since + last heartbeat), and ``seconds_since_heartbeat``. + """ + return self._client._get("/v1/workers") + + def delete(self, worker_id: str) -> None: + """Delete a registered worker. + + Removes the worker row; in-flight executions claimed by this + worker will be picked up by the stale-recovery loop. Useful for + cleaning up workers that have been decommissioned. + + Returns ``None`` on success (HTTP 204). Raises + ``CueNotFoundError`` if the worker doesn't exist. + + Args: + worker_id: The caller-defined worker_id used during the + worker's heartbeats. Same value is what appears in + ``list()`` responses. + """ + return self._client._delete(f"/v1/workers/{worker_id}") diff --git a/tests/test_usage_resource.py b/tests/test_usage_resource.py new file mode 100644 index 0000000..aa1be8b --- /dev/null +++ b/tests/test_usage_resource.py @@ -0,0 +1,34 @@ +"""Tests for UsageResource.""" + +from unittest.mock import MagicMock + +from cueapi.resources.usage import UsageResource + + +class TestGet: + def test_get_calls_get_usage(self): + mock_client = MagicMock() + mock_client._get.return_value = { + "plan": {"name": "pro", "interval": "monthly"}, + "cues": {"active": 12, "limit": 100}, + "executions": {"used": 543, "limit": 5000}, + "rate_limit": {"limit": 200}, + } + resource = UsageResource(mock_client) + + result = resource.get() + + mock_client._get.assert_called_once_with("/v1/usage") + assert result["plan"]["name"] == "pro" + assert result["cues"]["active"] == 12 + + def test_get_returns_server_dict_unchanged(self): + # Pin the no-transform behavior so a future refactor can't + # silently start coercing the response into a typed object + # without bumping the major version. + mock_client = MagicMock() + mock_client._get.return_value = {"unexpected_field": "value"} + resource = UsageResource(mock_client) + + result = resource.get() + assert result == {"unexpected_field": "value"} diff --git a/tests/test_workers_resource.py b/tests/test_workers_resource.py new file mode 100644 index 0000000..e9f96d3 --- /dev/null +++ b/tests/test_workers_resource.py @@ -0,0 +1,55 @@ +"""Tests for WorkersResource.""" + +from unittest.mock import MagicMock + +from cueapi.resources.workers import WorkersResource + + +class TestList: + def test_list_calls_get_workers(self): + mock_client = MagicMock() + mock_client._get.return_value = { + "workers": [ + { + "worker_id": "worker-1", + "handlers": ["task-a"], + "last_heartbeat": "2026-05-04T17:30:00Z", + "heartbeat_status": "online", + "seconds_since_heartbeat": 5, + } + ], + "total": 1, + } + resource = WorkersResource(mock_client) + + result = resource.list() + + mock_client._get.assert_called_once_with("/v1/workers") + assert result["total"] == 1 + assert result["workers"][0]["worker_id"] == "worker-1" + + def test_list_passes_no_params(self): + # Endpoint accepts no query params; SDK MUST NOT silently start + # passing params (would couple to a future server-side change). + # Pinning the bare-call shape. + mock_client = MagicMock() + mock_client._get.return_value = {"workers": [], "total": 0} + resource = WorkersResource(mock_client) + + resource.list() + + mock_client._get.assert_called_once_with("/v1/workers") + # No params kwarg. + assert "params" not in mock_client._get.call_args.kwargs + + +class TestDelete: + def test_delete_calls_delete_workers_id(self): + mock_client = MagicMock() + mock_client._delete.return_value = None # 204 -> None per client _request + resource = WorkersResource(mock_client) + + result = resource.delete("worker-1") + + mock_client._delete.assert_called_once_with("/v1/workers/worker-1") + assert result is None