From bf11e8f6952869d493942de701c441d591a55aef Mon Sep 17 00:00:00 2001 From: Poseidon Date: Sun, 12 Jul 2026 15:30:26 +0000 Subject: [PATCH 1/4] fix(job): surface review/spend-limited states and cost warning MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit RapidataJob._wait_for_status waited only for Completed/Failed, so a job that entered ManualApproval or SpendLimited blocked the caller forever (e.g. via get_results). It now raises an informative error naming the state — including the reviewReason when the API returns one, a generic "is being reviewed" message when it doesn't, and a partial-results / top-up hint for SpendLimited. display_progress_bar gets the same guard. assign_job now surfaces the create response's optional costWarning as a Python warning (estimated cost, available balance, shortfall), phrased as an estimate. The job is still created and behavior is otherwise unchanged. Adds pytest + tests covering both hang-fix states (with and without a reviewReason), the cost-warning emission, and the unchanged happy path, plus a docs note in audiences.md. Co-Authored-By: Claude Opus 4.8 Co-Authored-By: karl --- docs/audiences.md | 13 ++ pyproject.toml | 4 + .../audience/_audience_base.py | 25 +++ .../rapidata_client/job/rapidata_job.py | 55 ++++++- tests/test_job_review_and_cost_warning.py | 149 ++++++++++++++++++ uv.lock | 31 +++- 6 files changed, 275 insertions(+), 2 deletions(-) create mode 100644 tests/test_job_review_and_cost_warning.py diff --git a/docs/audiences.md b/docs/audiences.md index 21c967292..9b6b27100 100644 --- a/docs/audiences.md +++ b/docs/audiences.md @@ -99,6 +99,19 @@ results = job.get_results() print(results) ``` +### Cost warnings and jobs under review + +`assign_job` never blocks on funds: the job is always created. If its estimated cost +exceeds your account balance, `assign_job` emits a Python warning with the estimate, +your balance, and the expected shortfall — the job still runs, but may pause partway +until you top up. + +Some jobs don't go straight to running. A job can enter manual review +(`ManualApproval`) or, once out of funds mid-run, become spend-limited +(`SpendLimited`). Neither state completes on its own, so `get_results()` raises an +informative error naming the state (and the review reason, when available) instead of +blocking indefinitely — top up or wait for a reviewer, then call it again. + ## Complete Example Here's the full workflow — creating a custom audience, adding qualification examples, and running a labeling job: diff --git a/pyproject.toml b/pyproject.toml index 85327dec3..43c33c7dc 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -34,6 +34,7 @@ dev = [ "ipykernel>=7.1.0,<8", "jupyter>=1.1.1,<2", "pyright>=1.1.407", + "pytest>=8.0.0,<9", ] docs = [ "mkdocs-material>=9.5.34,<10", @@ -49,6 +50,9 @@ docs = [ "mkdocs-llmstxt-md>=0.1.0", ] +[tool.pytest.ini_options] +testpaths = ["tests"] + [tool.hatch.build.targets.sdist] include = ["src/rapidata"] diff --git a/src/rapidata/rapidata_client/audience/_audience_base.py b/src/rapidata/rapidata_client/audience/_audience_base.py index 87bc2a755..cd7a6fc9d 100644 --- a/src/rapidata/rapidata_client/audience/_audience_base.py +++ b/src/rapidata/rapidata_client/audience/_audience_base.py @@ -1,5 +1,6 @@ from __future__ import annotations +import warnings from typing import TYPE_CHECKING from rapidata.rapidata_client.config import logger, managed_print, tracer @@ -10,6 +11,9 @@ RapidataJobDefinition, ) from rapidata.rapidata_client.job.rapidata_job import RapidataJob + from rapidata.api_client.models.create_job_endpoint_cost_warning_model import ( + CreateJobEndpointCostWarningModel, + ) class RapidataAudienceBase: @@ -89,8 +93,29 @@ def assign_job(self, job_definition: RapidataJobDefinition) -> RapidataJob: managed_print( f"Job '{job.name}' is now viewable under: {job.job_details_page}" ) + self._warn_if_cost_exceeds_balance(job, response.cost_warning) return job + @staticmethod + def _warn_if_cost_exceeds_balance( + job: RapidataJob, cost_warning: CreateJobEndpointCostWarningModel | None + ) -> None: + """Surface the create response's optional cost warning as a Python warning. + + The job is created and runs regardless — this is an advisory estimate that it + may pause for funds before finishing, not an error. + """ + if cost_warning is None: + return + warnings.warn( + f"Job '{job.name}' has an estimated cost of {cost_warning.estimated_cost:.2f}, " + f"but the account balance is {cost_warning.available_balance:.2f} — it will " + f"likely pause about {cost_warning.shortfall:.2f} short of finishing. The job " + f"was created and runs as far as the balance allows; top up the account to let " + f"it complete. This is an estimate.", + stacklevel=3, + ) + def find_jobs( self, name: str = "", amount: int = 10, page: int = 1 ) -> list[RapidataJob]: diff --git a/src/rapidata/rapidata_client/job/rapidata_job.py b/src/rapidata/rapidata_client/job/rapidata_job.py index 129454f9b..f4db605d6 100644 --- a/src/rapidata/rapidata_client/job/rapidata_job.py +++ b/src/rapidata/rapidata_client/job/rapidata_job.py @@ -70,6 +70,12 @@ def __init__( self.job_details_page = f"https://app.{self._openapi_service.environment}/audiences/{self.audience_id}/job/{self.id}" logger.debug("RapidataJob initialized") + # States a job can settle into that will never progress to Completed/Failed on + # their own: ManualApproval needs a Rapidata reviewer to act, SpendLimited needs + # an account top-up. Waiting on either (e.g. from get_results) would hang the + # caller forever, so we surface them as informative errors instead. + _BLOCKING_STATUSES = ("ManualApproval", "SpendLimited") + def _get_job_failure_message(self) -> str | None: """Retrieves the failure message from the job if available.""" try: @@ -79,6 +85,36 @@ def _get_job_failure_message(self) -> str | None: logger.debug("Failed to get job failure message", self, exc_info=True) return None + def _get_review_reason(self) -> str | None: + """Retrieves the manual-review reason from the job if the API exposes one.""" + try: + job = self._openapi_service.order.job_api.job_job_id_get(self.id) + reason = job.review_reason + return reason.value if reason is not None else None + except Exception: + logger.debug("Failed to get job review reason", self, exc_info=True) + return None + + def _raise_for_blocking_status(self, status: str) -> None: + """Raises an informative error for a job state that can't reach completion + on its own, instead of letting the caller block on it indefinitely.""" + if status == "SpendLimited": + raise Exception( + f"Job '{self}' is spend-limited: the account ran out of funds while " + f"running, so it stopped collecting responses. Partial results remain " + f"available; top up the account to resume and let the job finish." + ) + + # ManualApproval — reviewReason is optional; a job can legitimately be under + # review with no customer-facing reason recorded yet. + reason = self._get_review_reason() + reason_detail = f" ({reason})" if reason else "" + raise Exception( + f"Job '{self}' is being reviewed{reason_detail}: a Rapidata reviewer " + f"needs to approve it before it runs. This can take a couple of hours; " + f"once approved, call this again." + ) + def _retry_operation( self, operation: Callable[[], T], @@ -133,8 +169,15 @@ def _wait_for_status( Returns: The final status reached + + Raises: + Exception: If the job enters a state it can't progress out of on its own + (``ManualApproval`` or ``SpendLimited``) while a different status is + being awaited — with the review reason when the API provides one. """ while (current_status := self.get_status()) not in target_statuses: + if current_status in self._BLOCKING_STATUSES: + self._raise_for_blocking_status(current_status) if status_message: logger.debug(status_message, self, current_status) sleep(check_interval) @@ -221,7 +264,9 @@ def get_results(self) -> RapidataResults: RapidataResults: The results of the job. Raises: - Exception: If failed to get job results. + Exception: If failed to get job results, or if the job is in manual + review (``ManualApproval``) or spend-limited (``SpendLimited``) and + therefore cannot complete without intervention. """ with tracer.start_as_current_span("RapidataJob.get_results"): from rapidata.api_client.exceptions import ApiException @@ -260,6 +305,8 @@ def display_progress_bar(self, refresh_rate: int = 5) -> None: Raises: ValueError: If refresh_rate is less than 1. + Exception: If the job has failed, or is in a state it can't progress out + of on its own (``ManualApproval`` or ``SpendLimited``). """ if refresh_rate < 1: raise ValueError("refresh_rate must be at least 1") @@ -273,6 +320,9 @@ def display_progress_bar(self, refresh_rate: int = 5) -> None: failure_message = self._get_job_failure_message() raise Exception(f"Job '{self}' has failed: {failure_message}") + if current_status in self._BLOCKING_STATUSES: + self._raise_for_blocking_status(current_status) + # Get progress from pipeline if available with tqdm( total=100, @@ -293,6 +343,9 @@ def display_progress_bar(self, refresh_rate: int = 5) -> None: failure_message = self._get_job_failure_message() raise Exception(f"Job '{self}' has failed: {failure_message}") + if current_status in self._BLOCKING_STATUSES: + self._raise_for_blocking_status(current_status) + # Try to get progress from workflow try: progress = self._get_workflow_progress() diff --git a/tests/test_job_review_and_cost_warning.py b/tests/test_job_review_and_cost_warning.py new file mode 100644 index 000000000..8fe37a904 --- /dev/null +++ b/tests/test_job_review_and_cost_warning.py @@ -0,0 +1,149 @@ +"""Tests for surfacing job review/spend-limited states and the create cost warning. + +Covers workstream E of the "cost warning + transparent review outcomes" work: +``RapidataJob`` must stop hanging on ``ManualApproval``/``SpendLimited`` and raise an +informative error instead, and ``assign_job`` must surface the create response's +optional ``costWarning`` as a Python warning without otherwise changing behavior. +""" + +from __future__ import annotations + +import json +from unittest.mock import MagicMock + +import pytest + +from rapidata.api_client.models.create_job_endpoint_cost_warning_model import ( + CreateJobEndpointCostWarningModel, +) +from rapidata.api_client.models.review_reason_model import ReviewReasonModel +from rapidata.rapidata_client.audience._audience_base import RapidataAudienceBase +from rapidata.rapidata_client.job.rapidata_job import RapidataJob + + +def _job_get(state: str, review_reason: ReviewReasonModel | None = None) -> MagicMock: + """A stand-in for the job GET response with just the fields the code reads.""" + job = MagicMock() + job.state.value = state + job.review_reason = review_reason + job.failure_message = None + return job + + +def _make_job(job_get: MagicMock) -> tuple[RapidataJob, MagicMock]: + openapi_service = MagicMock() + openapi_service.environment = "rapidata.ai" + openapi_service.order.job_api.job_job_id_get.return_value = job_get + job = RapidataJob( + job_id="job-1", + name="My Job", + audience_id="aud-1", + created_at=MagicMock(), + definition_id="def-1", + openapi_service=openapi_service, + ) + return job, openapi_service + + +def test_get_results_raises_on_manual_approval_with_reason(): + job, _ = _make_job( + _job_get("ManualApproval", review_reason=ReviewReasonModel.CONTENTFLAGGED) + ) + + with pytest.raises(Exception) as excinfo: + job.get_results() + + message = str(excinfo.value) + assert "being reviewed" in message + assert "ContentFlagged" in message + + +def test_get_results_raises_on_manual_approval_without_reason(): + job, _ = _make_job(_job_get("ManualApproval", review_reason=None)) + + with pytest.raises(Exception) as excinfo: + job.get_results() + + message = str(excinfo.value) + assert "being reviewed" in message + # A null reason must not leak into the message as empty parentheses. + assert "()" not in message + + +def test_get_results_raises_on_spend_limited(): + job, _ = _make_job(_job_get("SpendLimited")) + + with pytest.raises(Exception) as excinfo: + job.get_results() + + message = str(excinfo.value).lower() + assert "spend-limited" in message + assert "partial results" in message + assert "top up" in message + + +def test_get_results_happy_path_returns_results(): + job, openapi_service = _make_job(_job_get("Completed")) + openapi_service.order.job_api.job_job_id_download_results_get.return_value = ( + json.dumps({"info": {}, "results": []}) + ) + + results = job.get_results() + + assert results == {"info": {}, "results": []} + openapi_service.order.job_api.job_job_id_download_results_get.assert_called_once_with( + job_id="job-1" + ) + + +def test_display_progress_bar_raises_on_spend_limited(): + job, _ = _make_job(_job_get("SpendLimited")) + + with pytest.raises(Exception) as excinfo: + job.display_progress_bar() + + assert "spend-limited" in str(excinfo.value).lower() + + +def _make_audience() -> tuple[RapidataAudienceBase, MagicMock]: + openapi_service = MagicMock() + openapi_service.environment = "rapidata.ai" + audience = RapidataAudienceBase( + id="aud-1", name="Aud", filters=[], openapi_service=openapi_service + ) + return audience, openapi_service + + +def test_assign_job_warns_on_cost_warning(): + audience, openapi_service = _make_audience() + response = MagicMock() + response.job_id = "job-1" + response.cost_warning = CreateJobEndpointCostWarningModel( + estimatedCost=120.0, availableBalance=50.0, shortfall=70.0 + ) + openapi_service.order.job_api.job_post.return_value = response + + with pytest.warns(UserWarning) as record: + audience.assign_job(MagicMock(id="def-1", name="My Job")) + + message = str(record[0].message) + assert "120.00" in message + assert "50.00" in message + assert "70.00" in message + assert "estimate" in message.lower() + + +def test_assign_job_no_warning_when_cost_within_balance(): + audience, openapi_service = _make_audience() + response = MagicMock() + response.job_id = "job-1" + response.cost_warning = None + openapi_service.order.job_api.job_post.return_value = response + + import warnings as _warnings + + with _warnings.catch_warnings(): + _warnings.simplefilter("error") + job = audience.assign_job(MagicMock(id="def-1", name="My Job")) + + assert job.id == "job-1" diff --git a/uv.lock b/uv.lock index 0f823f568..b04fdaae5 100644 --- a/uv.lock +++ b/uv.lock @@ -799,6 +799,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/a4/ed/1f1afb2e9e7f38a545d628f864d562a5ae64fe6f7a10e28ffb9b185b4e89/importlib_resources-6.5.2-py3-none-any.whl", hash = "sha256:789cfdc3ed28c78b67a06acb8126751ced69a3d5f79c095a98298cd8a760ccec", size = 37461, upload-time = "2025-01-03T18:51:54.306Z" }, ] +[[package]] +name = "iniconfig" +version = "2.3.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/72/34/14ca021ce8e5dfedc35312d08ba8bf51fdd999c576889fc2c24cb97f4f10/iniconfig-2.3.0.tar.gz", hash = "sha256:c76315c77db068650d49c5b56314774a7804df16fee4402c1f19d6d15d8c4730", size = 20503, upload-time = "2025-10-18T21:55:43.219Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/cb/b1/3846dd7f199d53cb17f49cba7e651e9ce294d8497c8c150530ed11865bb8/iniconfig-2.3.0-py3-none-any.whl", hash = "sha256:f631c04d2c48c52b84d0d0549c99ff3859c98df65b3101406327ecc7d53fbf12", size = 7484, upload-time = "2025-10-18T21:55:41.639Z" }, +] + [[package]] name = "ipykernel" version = "7.1.0" @@ -2357,6 +2366,24 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/dc/93/b69052907d032b00c40cb656d21438ec00b3a471733de137a3f65a49a0a0/pyright-1.1.407-py3-none-any.whl", hash = "sha256:6dd419f54fcc13f03b52285796d65e639786373f433e243f8b94cf93a7444d21", size = 5997008, upload-time = "2025-10-24T23:17:13.159Z" }, ] +[[package]] +name = "pytest" +version = "8.4.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "colorama", marker = "sys_platform == 'win32'" }, + { name = "exceptiongroup", marker = "python_full_version < '3.11'" }, + { name = "iniconfig" }, + { name = "packaging" }, + { name = "pluggy" }, + { name = "pygments" }, + { name = "tomli", marker = "python_full_version < '3.11'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/a3/5c/00a0e072241553e1a7496d638deababa67c5058571567b92a7eaa258397c/pytest-8.4.2.tar.gz", hash = "sha256:86c0d0b93306b961d58d62a4db4879f27fe25513d4b969df351abdddb3c30e01", size = 1519618, upload-time = "2025-09-04T14:34:22.711Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a8/a4/20da314d277121d6534b3a980b29035dcd51e6744bd79075a6ce8fa4eb8d/pytest-8.4.2-py3-none-any.whl", hash = "sha256:872f880de3fc3a5bdc88a11b39c9710c3497a547cfa9320bc3c5e62fbf272e79", size = 365750, upload-time = "2025-09-04T14:34:20.226Z" }, +] + [[package]] name = "python-dateutil" version = "2.9.0.post0" @@ -2583,7 +2610,7 @@ wheels = [ [[package]] name = "rapidata" -version = "3.15.6" +version = "3.16.4" source = { editable = "." } dependencies = [ { name = "authlib" }, @@ -2612,6 +2639,7 @@ dev = [ { name = "ipykernel" }, { name = "jupyter" }, { name = "pyright" }, + { name = "pytest" }, { name = "python-dotenv" }, ] docs = [ @@ -2656,6 +2684,7 @@ dev = [ { name = "ipykernel", specifier = ">=7.1.0,<8" }, { name = "jupyter", specifier = ">=1.1.1,<2" }, { name = "pyright", specifier = ">=1.1.407" }, + { name = "pytest", specifier = ">=8.0.0,<9" }, { name = "python-dotenv", specifier = ">=1.0.1,<2" }, ] docs = [ From 535d8d617443118653a9edead38d8256a64dbd71 Mon Sep 17 00:00:00 2001 From: RapidPoseidon Date: Mon, 13 Jul 2026 09:58:27 +0000 Subject: [PATCH 2/4] refactor(job): fetch job once per poll and use the state enum MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Address code-review findings on the review-states work: - quickstart.md still promised get_results "blocks until the job is complete", which is no longer true for ManualApproval/SpendLimited — point at the new audiences.md section instead. - The polling loops fetched the job for its status, then fetched it again for the review reason (and failure message) when raising. They now fetch the GET output once per iteration and pass it to _raise_for_blocking_status; _get_review_reason is gone. - Blocking-state comparisons use the generated AudienceJobState enum instead of raw strings, matching how RapidataOrder uses OrderState. Co-Authored-By: Claude Fable 5 Co-Authored-By: karl --- docs/quickstart.md | 2 +- .../rapidata_client/job/rapidata_job.py | 74 +++++++++---------- tests/test_job_review_and_cost_warning.py | 3 +- 3 files changed, 39 insertions(+), 40 deletions(-) diff --git a/docs/quickstart.md b/docs/quickstart.md index f5c2a65fc..3ceba845b 100644 --- a/docs/quickstart.md +++ b/docs/quickstart.md @@ -105,7 +105,7 @@ results = job.get_results() # (3)! 1. Assigns the job definition to the audience and starts collecting responses. 2. Opens your browser on the running job, where you can watch responses come in and monitor progress. -3. Blocks until the job is complete and returns the results. You can also monitor progress on the [Rapidata Dashboard](https://app.rapidata.ai/dashboard). +3. Blocks until the job is complete and returns the results. If the job needs manual review or runs out of funds, this raises an informative error instead of blocking — see [Cost warnings and jobs under review](audiences.md#cost-warnings-and-jobs-under-review). You can also monitor progress on the [Rapidata Dashboard](https://app.rapidata.ai/dashboard). To understand the results format, see the [Understanding the Results](understanding_the_results.md) guide. diff --git a/src/rapidata/rapidata_client/job/rapidata_job.py b/src/rapidata/rapidata_client/job/rapidata_job.py index f4db605d6..b8ef61525 100644 --- a/src/rapidata/rapidata_client/job/rapidata_job.py +++ b/src/rapidata/rapidata_client/job/rapidata_job.py @@ -9,6 +9,7 @@ from colorama import Fore from tqdm.auto import tqdm +from rapidata.api_client.models.audience_job_state import AudienceJobState from rapidata.service.openapi_service import OpenAPIService from rapidata.rapidata_client.config import ( logger, @@ -27,6 +28,9 @@ ) if TYPE_CHECKING: + from rapidata.api_client.models.get_job_by_id_endpoint_output import ( + GetJobByIdEndpointOutput, + ) from rapidata.rapidata_client.results.rapidata_results import RapidataResults T = TypeVar("T") @@ -74,31 +78,27 @@ def __init__( # their own: ManualApproval needs a Rapidata reviewer to act, SpendLimited needs # an account top-up. Waiting on either (e.g. from get_results) would hang the # caller forever, so we surface them as informative errors instead. - _BLOCKING_STATUSES = ("ManualApproval", "SpendLimited") + _BLOCKING_STATUSES = ( + AudienceJobState.MANUALAPPROVAL, + AudienceJobState.SPENDLIMITED, + ) + + def _fetch_job(self) -> GetJobByIdEndpointOutput: + """Fetches the job's GET output (state, failure message, review reason).""" + return self._openapi_service.order.job_api.job_job_id_get(self.id) def _get_job_failure_message(self) -> str | None: """Retrieves the failure message from the job if available.""" try: - job = self._openapi_service.order.job_api.job_job_id_get(self.id) - return job.failure_message + return self._fetch_job().failure_message except Exception: logger.debug("Failed to get job failure message", self, exc_info=True) return None - def _get_review_reason(self) -> str | None: - """Retrieves the manual-review reason from the job if the API exposes one.""" - try: - job = self._openapi_service.order.job_api.job_job_id_get(self.id) - reason = job.review_reason - return reason.value if reason is not None else None - except Exception: - logger.debug("Failed to get job review reason", self, exc_info=True) - return None - - def _raise_for_blocking_status(self, status: str) -> None: + def _raise_for_blocking_status(self, job: GetJobByIdEndpointOutput) -> None: """Raises an informative error for a job state that can't reach completion on its own, instead of letting the caller block on it indefinitely.""" - if status == "SpendLimited": + if job.state == AudienceJobState.SPENDLIMITED: raise Exception( f"Job '{self}' is spend-limited: the account ran out of funds while " f"running, so it stopped collecting responses. Partial results remain " @@ -107,8 +107,8 @@ def _raise_for_blocking_status(self, status: str) -> None: # ManualApproval — reviewReason is optional; a job can legitimately be under # review with no customer-facing reason recorded yet. - reason = self._get_review_reason() - reason_detail = f" ({reason})" if reason else "" + reason = job.review_reason + reason_detail = f" ({reason.value})" if reason else "" raise Exception( f"Job '{self}' is being reviewed{reason_detail}: a Rapidata reviewer " f"needs to approve it before it runs. This can take a couple of hours; " @@ -175,15 +175,17 @@ def _wait_for_status( (``ManualApproval`` or ``SpendLimited``) while a different status is being awaited — with the review reason when the API provides one. """ - while (current_status := self.get_status()) not in target_statuses: + while True: + job = self._fetch_job() + current_status = job.state.value + if current_status in target_statuses: + return current_status if current_status in self._BLOCKING_STATUSES: - self._raise_for_blocking_status(current_status) + self._raise_for_blocking_status(job) if status_message: logger.debug(status_message, self, current_status) sleep(check_interval) - return current_status - @property def completed_at(self) -> datetime | None: """Returns the completion date of the job, or None if not completed.""" @@ -233,9 +235,7 @@ def get_status(self) -> str: The current status of the job as a string. """ with tracer.start_as_current_span("RapidataJob.get_status"): - return self._openapi_service.order.job_api.job_job_id_get( - self.id - ).state.value + return self._fetch_job().state.value def _regenerate_results(self) -> None: """Triggers regeneration of a job whose results have gone stale. @@ -311,17 +311,16 @@ def display_progress_bar(self, refresh_rate: int = 5) -> None: if refresh_rate < 1: raise ValueError("refresh_rate must be at least 1") - current_status = self.get_status() - if current_status == "Completed": + job = self._fetch_job() + if job.state == AudienceJobState.COMPLETED: managed_print(f"Job '{self}' is already completed.") return - if current_status == "Failed": - failure_message = self._get_job_failure_message() - raise Exception(f"Job '{self}' has failed: {failure_message}") + if job.state == AudienceJobState.FAILED: + raise Exception(f"Job '{self}' has failed: {job.failure_message}") - if current_status in self._BLOCKING_STATUSES: - self._raise_for_blocking_status(current_status) + if job.state in self._BLOCKING_STATUSES: + self._raise_for_blocking_status(job) # Get progress from pipeline if available with tqdm( @@ -333,18 +332,17 @@ def display_progress_bar(self, refresh_rate: int = 5) -> None: ) as pbar: last_percentage = 0 while True: - current_status = self.get_status() + job = self._fetch_job() - if current_status == "Completed": + if job.state == AudienceJobState.COMPLETED: pbar.update(100 - last_percentage) break - if current_status == "Failed": - failure_message = self._get_job_failure_message() - raise Exception(f"Job '{self}' has failed: {failure_message}") + if job.state == AudienceJobState.FAILED: + raise Exception(f"Job '{self}' has failed: {job.failure_message}") - if current_status in self._BLOCKING_STATUSES: - self._raise_for_blocking_status(current_status) + if job.state in self._BLOCKING_STATUSES: + self._raise_for_blocking_status(job) # Try to get progress from workflow try: diff --git a/tests/test_job_review_and_cost_warning.py b/tests/test_job_review_and_cost_warning.py index 8fe37a904..43fc1886b 100644 --- a/tests/test_job_review_and_cost_warning.py +++ b/tests/test_job_review_and_cost_warning.py @@ -13,6 +13,7 @@ import pytest +from rapidata.api_client.models.audience_job_state import AudienceJobState from rapidata.api_client.models.create_job_endpoint_cost_warning_model import ( CreateJobEndpointCostWarningModel, ) @@ -24,7 +25,7 @@ def _job_get(state: str, review_reason: ReviewReasonModel | None = None) -> MagicMock: """A stand-in for the job GET response with just the fields the code reads.""" job = MagicMock() - job.state.value = state + job.state = AudienceJobState(state) job.review_reason = review_reason job.failure_message = None return job From 7ffa27020f0aada54c5b8c5716ba49f19eaf45a7 Mon Sep 17 00:00:00 2001 From: Poseidon Date: Mon, 13 Jul 2026 12:00:01 +0000 Subject: [PATCH 3/4] refactor(job): use SDK logger for cost warning; restructure tests Address review feedback from @LinoGiger: - Emit the create-time cost warning via the SDK logger (logger.warning, the repo's user-facing warning convention) instead of warnings.warn. - Move tests into a package layout mirroring src (tests/rapidata_client/{job,audience}/) instead of one flat file, and keep the AudienceJobState-based test helper from the polling refactor. Co-Authored-By: Claude Opus 4.8 Co-Authored-By: karl --- docs/audiences.md | 6 +- .../audience/_audience_base.py | 19 ++--- tests/rapidata_client/__init__.py | 0 tests/rapidata_client/audience/__init__.py | 0 .../audience/test_audience_base.py | 64 ++++++++++++++++ tests/rapidata_client/job/__init__.py | 0 .../job/test_rapidata_job.py} | 73 +++---------------- 7 files changed, 89 insertions(+), 73 deletions(-) create mode 100644 tests/rapidata_client/__init__.py create mode 100644 tests/rapidata_client/audience/__init__.py create mode 100644 tests/rapidata_client/audience/test_audience_base.py create mode 100644 tests/rapidata_client/job/__init__.py rename tests/{test_job_review_and_cost_warning.py => rapidata_client/job/test_rapidata_job.py} (58%) diff --git a/docs/audiences.md b/docs/audiences.md index 9b6b27100..3723afb1c 100644 --- a/docs/audiences.md +++ b/docs/audiences.md @@ -102,9 +102,9 @@ print(results) ### Cost warnings and jobs under review `assign_job` never blocks on funds: the job is always created. If its estimated cost -exceeds your account balance, `assign_job` emits a Python warning with the estimate, -your balance, and the expected shortfall — the job still runs, but may pause partway -until you top up. +exceeds your account balance, `assign_job` logs a warning with the estimate, your +balance, and the expected shortfall — the job still runs, but may pause partway until +you top up. Some jobs don't go straight to running. A job can enter manual review (`ManualApproval`) or, once out of funds mid-run, become spend-limited diff --git a/src/rapidata/rapidata_client/audience/_audience_base.py b/src/rapidata/rapidata_client/audience/_audience_base.py index cd7a6fc9d..1ae0735be 100644 --- a/src/rapidata/rapidata_client/audience/_audience_base.py +++ b/src/rapidata/rapidata_client/audience/_audience_base.py @@ -1,6 +1,5 @@ from __future__ import annotations -import warnings from typing import TYPE_CHECKING from rapidata.rapidata_client.config import logger, managed_print, tracer @@ -100,20 +99,22 @@ def assign_job(self, job_definition: RapidataJobDefinition) -> RapidataJob: def _warn_if_cost_exceeds_balance( job: RapidataJob, cost_warning: CreateJobEndpointCostWarningModel | None ) -> None: - """Surface the create response's optional cost warning as a Python warning. + """Surface the create response's optional cost warning via the SDK logger. The job is created and runs regardless — this is an advisory estimate that it may pause for funds before finishing, not an error. """ if cost_warning is None: return - warnings.warn( - f"Job '{job.name}' has an estimated cost of {cost_warning.estimated_cost:.2f}, " - f"but the account balance is {cost_warning.available_balance:.2f} — it will " - f"likely pause about {cost_warning.shortfall:.2f} short of finishing. The job " - f"was created and runs as far as the balance allows; top up the account to let " - f"it complete. This is an estimate.", - stacklevel=3, + logger.warning( + "Job '%s' has an estimated cost of %.2f, but the account balance is %.2f — " + "it will likely pause about %.2f short of finishing. The job was created and " + "runs as far as the balance allows; top up the account to let it complete. " + "This is an estimate.", + job.name, + cost_warning.estimated_cost, + cost_warning.available_balance, + cost_warning.shortfall, ) def find_jobs( diff --git a/tests/rapidata_client/__init__.py b/tests/rapidata_client/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/rapidata_client/audience/__init__.py b/tests/rapidata_client/audience/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/rapidata_client/audience/test_audience_base.py b/tests/rapidata_client/audience/test_audience_base.py new file mode 100644 index 000000000..1256418a1 --- /dev/null +++ b/tests/rapidata_client/audience/test_audience_base.py @@ -0,0 +1,64 @@ +"""Tests for the create-time cost warning surfaced by assign_job. + +When job creation returns a costWarning (estimate exceeds balance), assign_job +logs an advisory warning. The job is created regardless; behaviour is otherwise +unchanged. +""" + +from __future__ import annotations + +from unittest.mock import MagicMock + +import rapidata.rapidata_client.audience._audience_base as base_module +from rapidata.api_client.models.create_job_endpoint_cost_warning_model import ( + CreateJobEndpointCostWarningModel, +) +from rapidata.rapidata_client.audience._audience_base import RapidataAudienceBase + + +def _make_audience() -> tuple[RapidataAudienceBase, MagicMock]: + openapi_service = MagicMock() + openapi_service.environment = "rapidata.ai" + audience = RapidataAudienceBase( + id="aud-1", name="Aud", filters=[], openapi_service=openapi_service + ) + return audience, openapi_service + + +def test_assign_job_warns_on_cost_warning(monkeypatch): + audience, openapi_service = _make_audience() + response = MagicMock() + response.job_id = "job-1" + response.cost_warning = CreateJobEndpointCostWarningModel( + estimatedCost=120.0, availableBalance=50.0, shortfall=70.0 + ) + openapi_service.order.job_api.job_post.return_value = response + + warn = MagicMock() + monkeypatch.setattr(base_module.logger, "warning", warn) + + audience.assign_job(MagicMock(id="def-1", name="My Job")) + + warn.assert_called_once() + args = warn.call_args.args + assert "estimate" in args[0].lower() + # Values are passed as lazy %-args, not pre-formatted into the message. + assert 120.0 in args + assert 50.0 in args + assert 70.0 in args + + +def test_assign_job_no_warning_when_cost_within_balance(monkeypatch): + audience, openapi_service = _make_audience() + response = MagicMock() + response.job_id = "job-1" + response.cost_warning = None + openapi_service.order.job_api.job_post.return_value = response + + warn = MagicMock() + monkeypatch.setattr(base_module.logger, "warning", warn) + + job = audience.assign_job(MagicMock(id="def-1", name="My Job")) + + warn.assert_not_called() + assert job.id == "job-1" diff --git a/tests/rapidata_client/job/__init__.py b/tests/rapidata_client/job/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/test_job_review_and_cost_warning.py b/tests/rapidata_client/job/test_rapidata_job.py similarity index 58% rename from tests/test_job_review_and_cost_warning.py rename to tests/rapidata_client/job/test_rapidata_job.py index 43fc1886b..6512e4ff8 100644 --- a/tests/test_job_review_and_cost_warning.py +++ b/tests/rapidata_client/job/test_rapidata_job.py @@ -1,9 +1,8 @@ -"""Tests for surfacing job review/spend-limited states and the create cost warning. +"""Tests for RapidataJob surfacing states it can't progress out of on its own. -Covers workstream E of the "cost warning + transparent review outcomes" work: -``RapidataJob`` must stop hanging on ``ManualApproval``/``SpendLimited`` and raise an -informative error instead, and ``assign_job`` must surface the create response's -optional ``costWarning`` as a Python warning without otherwise changing behavior. +A job that enters ManualApproval or SpendLimited never reaches Completed/Failed, +so waiting on it (e.g. via get_results) must raise an informative error instead +of blocking the caller forever. """ from __future__ import annotations @@ -14,11 +13,7 @@ import pytest from rapidata.api_client.models.audience_job_state import AudienceJobState -from rapidata.api_client.models.create_job_endpoint_cost_warning_model import ( - CreateJobEndpointCostWarningModel, -) from rapidata.api_client.models.review_reason_model import ReviewReasonModel -from rapidata.rapidata_client.audience._audience_base import RapidataAudienceBase from rapidata.rapidata_client.job.rapidata_job import RapidataJob @@ -83,20 +78,6 @@ def test_get_results_raises_on_spend_limited(): assert "top up" in message -def test_get_results_happy_path_returns_results(): - job, openapi_service = _make_job(_job_get("Completed")) - openapi_service.order.job_api.job_job_id_download_results_get.return_value = ( - json.dumps({"info": {}, "results": []}) - ) - - results = job.get_results() - - assert results == {"info": {}, "results": []} - openapi_service.order.job_api.job_job_id_download_results_get.assert_called_once_with( - job_id="job-1" - ) - - def test_display_progress_bar_raises_on_spend_limited(): job, _ = _make_job(_job_get("SpendLimited")) @@ -106,45 +87,15 @@ def test_display_progress_bar_raises_on_spend_limited(): assert "spend-limited" in str(excinfo.value).lower() -def _make_audience() -> tuple[RapidataAudienceBase, MagicMock]: - openapi_service = MagicMock() - openapi_service.environment = "rapidata.ai" - audience = RapidataAudienceBase( - id="aud-1", name="Aud", filters=[], openapi_service=openapi_service +def test_get_results_happy_path_returns_results(): + job, openapi_service = _make_job(_job_get("Completed")) + openapi_service.order.job_api.job_job_id_download_results_get.return_value = ( + json.dumps({"info": {}, "results": []}) ) - return audience, openapi_service + results = job.get_results() -def test_assign_job_warns_on_cost_warning(): - audience, openapi_service = _make_audience() - response = MagicMock() - response.job_id = "job-1" - response.cost_warning = CreateJobEndpointCostWarningModel( - estimatedCost=120.0, availableBalance=50.0, shortfall=70.0 + assert results == {"info": {}, "results": []} + openapi_service.order.job_api.job_job_id_download_results_get.assert_called_once_with( + job_id="job-1" ) - openapi_service.order.job_api.job_post.return_value = response - - with pytest.warns(UserWarning) as record: - audience.assign_job(MagicMock(id="def-1", name="My Job")) - - message = str(record[0].message) - assert "120.00" in message - assert "50.00" in message - assert "70.00" in message - assert "estimate" in message.lower() - - -def test_assign_job_no_warning_when_cost_within_balance(): - audience, openapi_service = _make_audience() - response = MagicMock() - response.job_id = "job-1" - response.cost_warning = None - openapi_service.order.job_api.job_post.return_value = response - - import warnings as _warnings - - with _warnings.catch_warnings(): - _warnings.simplefilter("error") - job = audience.assign_job(MagicMock(id="def-1", name="My Job")) - - assert job.id == "job-1" From 002c837d39c0799c718ef4e11922038216a07742 Mon Sep 17 00:00:00 2001 From: Lino Giger <68745352+LinoGiger@users.noreply.github.com> Date: Mon, 13 Jul 2026 14:38:20 +0200 Subject: [PATCH 4/4] added tests settings --- .vscode/settings.json | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/.vscode/settings.json b/.vscode/settings.json index a6c7dace8..752653713 100644 --- a/.vscode/settings.json +++ b/.vscode/settings.json @@ -6,8 +6,8 @@ "-p", "test_*.py" ], - "python.testing.unittestEnabled": true, - "python.testing.pytestEnabled": false, + "python.testing.unittestEnabled": false, + "python.testing.pytestEnabled": true, "python.autoComplete.extraPaths": [ "./src" ], @@ -19,5 +19,8 @@ "cursorpyright.analysis.extraPaths": [ "./src" ], - "cursorpyright.analysis.typeCheckingMode": "basic" + "cursorpyright.analysis.typeCheckingMode": "basic", + "python.testing.pytestArgs": [ + "tests" + ] }