diff --git a/plugins/nemo-data-designer/src/nemo_data_designer_plugin/sdk/errors.py b/plugins/nemo-data-designer/src/nemo_data_designer_plugin/sdk/errors.py index 9707ea8e6d..22c587ca0a 100644 --- a/plugins/nemo-data-designer/src/nemo_data_designer_plugin/sdk/errors.py +++ b/plugins/nemo-data-designer/src/nemo_data_designer_plugin/sdk/errors.py @@ -1,11 +1,25 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +from __future__ import annotations + +from contextlib import suppress + +import httpx from data_designer.errors import DataDesignerError class DataDesignerClientError(DataDesignerError): - """Base exception for Data Designer client errors.""" + """Base exception for Data Designer client errors. + + When the error originated from an HTTP response, the underlying status code + is exposed as :attr:`status_code` so callers can branch on it cleanly + instead of pattern-matching the message string. + """ + + def __init__(self, *args: object, status_code: int | None = None) -> None: + super().__init__(*args) + self.status_code = status_code class DataDesignerConfigValidationError(DataDesignerClientError): @@ -18,3 +32,26 @@ class DataDesignerPreviewError(DataDesignerClientError): class DataDesignerJobError(DataDesignerClientError): """Raised for errors related to a Data Designer job.""" + + +def extract_http_error_info(exc: httpx.HTTPStatusError) -> tuple[int, str]: + """Pull the status code and a human-readable detail string out of an httpx error. + + Tries to parse the response body as JSON and use its ``detail`` field (the + convention used by FastAPI / NeMo Platform); falls back to the raw body + text if that isn't available. + """ + response = exc.response + try: + response.read() + except Exception: + pass + + detail = response.text + body = None + with suppress(Exception): + body = response.json() + if isinstance(body, dict) and isinstance(body.get("detail"), str): + detail = body["detail"] + + return response.status_code, detail diff --git a/plugins/nemo-data-designer/src/nemo_data_designer_plugin/sdk/job_resources.py b/plugins/nemo-data-designer/src/nemo_data_designer_plugin/sdk/job_resources.py index e348561966..d30150edc0 100644 --- a/plugins/nemo-data-designer/src/nemo_data_designer_plugin/sdk/job_resources.py +++ b/plugins/nemo-data-designer/src/nemo_data_designer_plugin/sdk/job_resources.py @@ -16,7 +16,7 @@ from data_designer.config.utils.visualization import WithRecordSamplerMixin from data_designer.logging import RandomEmoji from nemo_data_designer_plugin.sdk import http -from nemo_data_designer_plugin.sdk.errors import DataDesignerJobError +from nemo_data_designer_plugin.sdk.errors import DataDesignerJobError, extract_http_error_info from nemo_data_designer_plugin.sdk.job_results import DataDesignerJobResults from nemo_data_designer_plugin.sdk.logging import with_logging from nemo_platform import AsyncNeMoPlatform, NeMoPlatform @@ -54,8 +54,8 @@ def _raise_for_status(resp: httpx.Response) -> None: try: resp.raise_for_status() except httpx.HTTPStatusError as exc: - detail = exc.response.text - raise DataDesignerJobError(detail) from exc + status_code, detail = extract_http_error_info(exc) + raise DataDesignerJobError(detail, status_code=status_code) from exc @dataclass @@ -271,7 +271,7 @@ def _check_if_result_available(self, result_name: str) -> None: else: logger.warning(f"Job ended with status {status!r}. Fetching completed {result_name} result.") except DataDesignerJobError as e: - if "404" in str(e): + if e.status_code == 404: raise DataDesignerJobError(f"{result_name!r} result is not available.") from e raise DataDesignerJobError(f"🛑 Error loading dataset: {e}") from e else: @@ -478,7 +478,7 @@ async def _check_if_result_available(self, result_name: str) -> None: else: logger.warning(f"Job ended with status {status!r}. Fetching completed {result_name} result.") except DataDesignerJobError as e: - if "404" in str(e): + if e.status_code == 404: raise DataDesignerJobError(f"{result_name!r} result is not available.") from e raise DataDesignerJobError(f"🛑 Error loading dataset: {e}") from e else: diff --git a/plugins/nemo-data-designer/src/nemo_data_designer_plugin/sdk/resources.py b/plugins/nemo-data-designer/src/nemo_data_designer_plugin/sdk/resources.py index 4ad36db5db..f956a3e791 100644 --- a/plugins/nemo-data-designer/src/nemo_data_designer_plugin/sdk/resources.py +++ b/plugins/nemo-data-designer/src/nemo_data_designer_plugin/sdk/resources.py @@ -33,6 +33,7 @@ DataDesignerClientError, DataDesignerConfigValidationError, DataDesignerPreviewError, + extract_http_error_info, ) from nemo_data_designer_plugin.sdk.job_resources import AsyncDataDesignerJobResource, DataDesignerJobResource from nemo_data_designer_plugin.sdk.logging import with_logging @@ -505,20 +506,10 @@ def _get_config_for_api_call(config_builder: dd.DataDesignerConfigBuilder) -> dd def _get_error(e: BaseException) -> DataDesignerClientError: if isinstance(e, httpx.HTTPStatusError): - try: - e.response.read() - except Exception: - pass - - detail = e.response.text - try: - detail_json = e.response.json() - detail = detail_json.get("detail", detail) - except Exception: - pass - if e.response.status_code == 422: - return DataDesignerConfigValidationError(f"‼️ Config validation failed!\n{detail}") - return DataDesignerClientError(f"‼️ Something went wrong!\n{detail}") + status_code, detail = extract_http_error_info(e) + if status_code == 422: + return DataDesignerConfigValidationError(f"‼️ Config validation failed!\n{detail}", status_code=status_code) + return DataDesignerClientError(f"‼️ Something went wrong!\n{detail}", status_code=status_code) return DataDesignerClientError(f"‼️ Something went wrong!\n{e}") diff --git a/plugins/nemo-data-designer/src/nemo_data_designer_plugin/testing/utils.py b/plugins/nemo-data-designer/src/nemo_data_designer_plugin/testing/utils.py index 0dc9578ab7..4a8722e115 100644 --- a/plugins/nemo-data-designer/src/nemo_data_designer_plugin/testing/utils.py +++ b/plugins/nemo-data-designer/src/nemo_data_designer_plugin/testing/utils.py @@ -10,14 +10,16 @@ from contextlib import asynccontextmanager, contextmanager, redirect_stderr, redirect_stdout from dataclasses import dataclass from pathlib import Path -from typing import Any +from typing import Any, Literal from unittest.mock import AsyncMock, Mock, patch from urllib.parse import unquote, urlparse +import click.testing import data_designer.config as dd import duckdb import pandas as pd import typer +import typer.testing from data_designer.engine.resources.seed_reader import SeedReader from data_designer_nemo.nemotron_personas import get_file_path_for_locale, get_resource_name_for_locale from nemo_data_designer_plugin.cli.main import DataDesignerCLI @@ -252,7 +254,7 @@ def make_dd_client(client_context: ClientContext) -> DataDesignerResource: return DataDesignerResource(client_context.sdk) -def make_data_designer_cli_app() -> typer.Typer: +def _make_data_designer_cli_app() -> typer.Typer: cli = DataDesignerCLI() app = cli.get_cli() add_function_commands(app, {"preview": PreviewFunction}, cli=cli) @@ -273,11 +275,13 @@ def get_async_client(self) -> AsyncNeMoPlatform: return self.async_sdk -def make_data_designer_cli_state( +def _make_data_designer_cli_state( client_context: ClientContext, *, output_format: str | None = None, ) -> DataDesignerCLIState: + # Mirrors what `nemo --output-format json` would do at the top-level callback. + # The plugin's test app doesn't mount the real top-level callback, so we set this directly. overrides = {"output_format": output_format} if output_format is not None else {} return DataDesignerCLIState( sdk=client_context.sdk, @@ -286,6 +290,21 @@ def make_data_designer_cli_state( ) +def invoke_cli( + command: list[str], + client_context: ClientContext | None = None, + output_format: Literal["json"] | None = None, +) -> click.testing.Result: + runner = typer.testing.CliRunner() + app = _make_data_designer_cli_app() + + cli_state = None + if client_context is not None: + cli_state = _make_data_designer_cli_state(client_context, output_format=output_format) + + return runner.invoke(app, command, obj=cli_state) + + def write_config_file(tmp_path: Path, source: str, *, name: str = "data_designer_config.py") -> Path: path = tmp_path / name path.write_text(source, encoding="utf-8") @@ -405,7 +424,10 @@ def __init__(self) -> None: workspace="default", name=job_name, source="data-designer", - spec={}, + # Store the canonical DataDesignerStepConfig as the job's spec so that + # downstream Data Designer routes (e.g. ``GET /jobs/create/{name}``, + # which deserializes the stored spec back through the schema) succeed. + spec=step_config, platform_spec=job_config_dict, ) job_ctx = JobContext( diff --git a/plugins/nemo-data-designer/tests/integration/test_job_sdk.py b/plugins/nemo-data-designer/tests/integration/test_job_sdk.py new file mode 100644 index 0000000000..ca74a22691 --- /dev/null +++ b/plugins/nemo-data-designer/tests/integration/test_job_sdk.py @@ -0,0 +1,347 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Integration tests for the Data Designer job-resource lifecycle. + +These tests exercise the end-user flow of inspecting a Data Designer job's +status/logs and downloading its artifacts — all the way through the in-process +FastAPI app, the Jobs core service, and the local file backend. + +A note on ``get_job_status``: ``task_context`` invokes ``CreateJob.run(...)`` +directly, which writes results and emits logs but does **not** drive the Jobs +service controller that would normally roll a step's terminal status up to +the job's ``completed`` state. We patch ``get_job_status`` on the resource as +the (single) network seam to bridge that gap. Everything past that seam — the +``_WaitLogCollector`` filtering, ``_status_is_complete`` branching, log-level +routing, the result download tarball extraction — runs unmocked against the +real services. +""" + +from collections.abc import AsyncGenerator, Generator +from contextlib import asynccontextmanager, contextmanager +from pathlib import Path +from unittest.mock import AsyncMock, patch + +import data_designer.config as dd +import nemo_data_designer_plugin.testing.utils as u +import pandas as pd +import pytest +from data_designer.config.analysis.dataset_profiler import DatasetProfilerResults +from nemo_data_designer_plugin.jobs.spec import DataDesignerJobConfig +from nemo_data_designer_plugin.sdk.errors import DataDesignerClientError, DataDesignerJobError +from nemo_data_designer_plugin.sdk.job_resources import ( + AsyncDataDesignerJobResource, + DataDesignerJobResource, +) +from nemo_data_designer_plugin.sdk.job_results import DataDesignerJobResults +from nemo_data_designer_plugin.sdk.resources import AsyncDataDesignerResource, DataDesignerResource + +pytestmark = pytest.mark.integration + +_JOB_NAME = "data-designer-abc123" + + +def _make_basic_job_config() -> DataDesignerJobConfig: + builder = dd.DataDesignerConfigBuilder(model_configs=[u.make_model_config()]) + builder.add_column( + column_config=dd.SamplerColumnConfig( + name="foo", + sampler_type=dd.SamplerType.CATEGORY, + params=dd.CategorySamplerParams(values=["a"]), + ) + ) + return DataDesignerJobConfig(num_records=3, config=builder.build()) + + +@asynccontextmanager +async def _completed_job() -> AsyncGenerator[u.CreateJobTestContext]: + """Stand up a real job and run it to completion (writes results, emits logs).""" + job_config = await u.compile_create_job(_make_basic_job_config(), workspace="default") + async with u.task_context(job_config, _JOB_NAME) as ctx: + result = ctx.run_task() + assert result.exit_code == 0, "task did not complete successfully" + yield ctx + + +@asynccontextmanager +async def _pending_job() -> AsyncGenerator[u.CreateJobTestContext]: + """Stand up a real job without running it (no results populated).""" + job_config = await u.compile_create_job(_make_basic_job_config(), workspace="default") + async with u.task_context(job_config, _JOB_NAME) as ctx: + yield ctx + + +@contextmanager +def _patch_status(resource: DataDesignerJobResource, status: str) -> Generator[None]: + """Mock the resource's network status call to return a specific platform-job status.""" + with patch.object(resource, "get_job_status", return_value=status): + yield + + +@contextmanager +def _patch_async_status(resource: AsyncDataDesignerJobResource, status: str) -> Generator[None]: + with patch.object(resource, "get_job_status", new=AsyncMock(return_value=status)): + yield + + +@contextmanager +def _no_pause() -> Generator[None]: + """Skip ``time.sleep`` calls inside ``wait_until_done`` so tests stay fast.""" + with ( + patch("nemo_data_designer_plugin.sdk.job_resources._pause"), + patch("nemo_data_designer_plugin.sdk.job_resources._async_pause"), + ): + yield + + +# --------------------------------------------------------------------------- +# get_job_resource / get_job +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_get_job_resource_returns_job_for_real_job() -> None: + async with _completed_job() as ctx: + dd_client = DataDesignerResource(ctx.sdk) + job_resource = dd_client.get_job_resource(_JOB_NAME, workspace="default") + + assert isinstance(job_resource, DataDesignerJobResource) + + job = job_resource.get_job() + assert job["name"] == _JOB_NAME + + +@pytest.mark.asyncio +async def test_get_job_resource_async_returns_job_for_real_job() -> None: + async with _completed_job() as ctx: + dd_client = AsyncDataDesignerResource(ctx.async_sdk) + job_resource = await dd_client.get_job_resource(_JOB_NAME, workspace="default") + + assert isinstance(job_resource, AsyncDataDesignerJobResource) + + job = await job_resource.get_job() + assert job["name"] == _JOB_NAME + + +# --------------------------------------------------------------------------- +# check_if_complete / _status_is_complete +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_check_if_complete_returns_true_for_completed_status() -> None: + async with _pending_job() as ctx: + job_resource = DataDesignerResource(ctx.sdk).get_job_resource(_JOB_NAME, workspace="default") + with _patch_status(job_resource, "completed"): + assert job_resource.check_if_complete() is True + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("simulated_status", "expected_log_fragment"), + [ + ("active", "still running"), + ("created", "still in the queue"), + ("pending", "still in the queue"), + ("error", "stopped with status `error`"), + ("cancelled", "stopped with status `cancelled`"), + ("frobnicated", "unknown state"), + ], +) +async def test_check_if_complete_returns_false_with_friendly_message_for_non_completed_statuses( + simulated_status: str, + expected_log_fragment: str, + caplog: pytest.LogCaptureFixture, +) -> None: + """Non-raising path should log a user-friendly message for every non-completed status.""" + + async with _pending_job() as ctx: + job_resource = DataDesignerResource(ctx.sdk).get_job_resource(_JOB_NAME, workspace="default") + + with _patch_status(job_resource, simulated_status), caplog.at_level("WARNING"): + assert job_resource.check_if_complete(raise_if_not_complete=False) is False + + assert any(expected_log_fragment in record.message for record in caplog.records), ( + f"expected log message containing {expected_log_fragment!r}, got: {[r.message for r in caplog.records]}" + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("simulated_status", ["active", "created", "pending", "error", "cancelled", "frobnicated"]) +async def test_check_if_complete_raises_when_requested(simulated_status: str) -> None: + async with _pending_job() as ctx: + job_resource = DataDesignerResource(ctx.sdk).get_job_resource(_JOB_NAME, workspace="default") + + with _patch_status(job_resource, simulated_status): + with pytest.raises(DataDesignerJobError): + job_resource.check_if_complete(raise_if_not_complete=True) + + +# --------------------------------------------------------------------------- +# wait_until_done +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_wait_until_done_logs_success_when_status_completes(caplog: pytest.LogCaptureFixture) -> None: + async with _completed_job() as ctx: + job_resource = DataDesignerResource(ctx.sdk).get_job_resource(_JOB_NAME, workspace="default") + with _no_pause(), _patch_status(job_resource, "completed"), caplog.at_level("INFO"): + job_resource.wait_until_done() + + assert any("completed successfully" in record.message for record in caplog.records) + + +@pytest.mark.asyncio +async def test_wait_until_done_async_logs_success_when_status_completes(caplog: pytest.LogCaptureFixture) -> None: + async with _completed_job() as ctx: + async_dd_client = AsyncDataDesignerResource(ctx.async_sdk) + job_resource = await async_dd_client.get_job_resource(_JOB_NAME, workspace="default") + with _no_pause(), _patch_async_status(job_resource, "completed"), caplog.at_level("INFO"): + await job_resource.wait_until_done() + + assert any("completed successfully" in record.message for record in caplog.records) + + +@pytest.mark.asyncio +async def test_wait_until_done_logs_terminal_failure_for_cancelled_status( + caplog: pytest.LogCaptureFixture, +) -> None: + async with _pending_job() as ctx: + job_resource = DataDesignerResource(ctx.sdk).get_job_resource(_JOB_NAME, workspace="default") + + with _no_pause(), _patch_status(job_resource, "cancelled"), caplog.at_level("ERROR"): + job_resource.wait_until_done() + + assert any("Terminating generation job" in record.message for record in caplog.records) + assert any("cancelled" in record.message for record in caplog.records) + + +# --------------------------------------------------------------------------- +# get_logs +# +# Note: ``DataDesignerJobResource.get_logs`` paginates through Job logs returned +# by the Files service's OTLP endpoint. ``task_context`` runs ``CreateJob.run`` +# in-process and bypasses the OTLP log-capture pipeline a real container runner +# would populate, so ``get_logs`` always returns ``[]`` here regardless of what +# the task emitted. +# --------------------------------------------------------------------------- + + +# --------------------------------------------------------------------------- +# download_artifacts (sync + async) and DataDesignerJobResults +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_download_artifacts_extracts_dataset_and_loads_analysis(tmp_path: Path) -> None: + async with _completed_job() as ctx: + job_resource = DataDesignerResource(ctx.sdk).get_job_resource(_JOB_NAME, workspace="default") + with _patch_status(job_resource, "completed"): + results = job_resource.download_artifacts(tmp_path) + + assert isinstance(results, DataDesignerJobResults) + + dataset = results.load_dataset() + assert isinstance(dataset, pd.DataFrame) + assert len(dataset) == 3 + assert dataset["foo"].tolist() == ["a", "a", "a"] + + analysis = results.load_analysis() + assert isinstance(analysis, DatasetProfilerResults) + assert analysis.num_records == 3 + + +@pytest.mark.asyncio +async def test_download_artifacts_async_extracts_dataset_and_loads_analysis(tmp_path: Path) -> None: + async with _completed_job() as ctx: + async_dd_client = AsyncDataDesignerResource(ctx.async_sdk) + job_resource = await async_dd_client.get_job_resource(_JOB_NAME, workspace="default") + with _patch_async_status(job_resource, "completed"): + results = await job_resource.download_artifacts(tmp_path) + + assert isinstance(results, DataDesignerJobResults) + assert results.load_analysis().num_records == 3 + + +@pytest.mark.asyncio +async def test_load_processor_dataset_raises_for_unknown_processor(tmp_path: Path) -> None: + async with _completed_job() as ctx: + job_resource = DataDesignerResource(ctx.sdk).get_job_resource(_JOB_NAME, workspace="default") + with _patch_status(job_resource, "completed"): + results = job_resource.download_artifacts(tmp_path) + + with pytest.raises(DataDesignerClientError, match="No artifacts found for processor"): + results.load_processor_dataset("undefined-processor") + + +# --------------------------------------------------------------------------- +# load_analysis (resource-level) and _check_if_result_available +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_load_analysis_returns_profiler_results_for_completed_status() -> None: + async with _completed_job() as ctx: + job_resource = DataDesignerResource(ctx.sdk).get_job_resource(_JOB_NAME, workspace="default") + with _patch_status(job_resource, "completed"): + analysis = job_resource.load_analysis() + + assert isinstance(analysis, DatasetProfilerResults) + assert analysis.num_records == 3 + + +@pytest.mark.asyncio +async def test_load_analysis_raises_when_status_is_unknown() -> None: + async with _pending_job() as ctx: + job_resource = DataDesignerResource(ctx.sdk).get_job_resource(_JOB_NAME, workspace="default") + + with _patch_status(job_resource, "frobnicated"): + with pytest.raises(DataDesignerJobError, match="frobnicated"): + job_resource.load_analysis() + + +@pytest.mark.asyncio +async def test_load_analysis_when_active_uses_completed_result_if_available( + caplog: pytest.LogCaptureFixture, +) -> None: + """``_check_if_result_available`` allows fetching completed results from an ``active`` job + and emits a 'still cooking' info message. + """ + async with _completed_job() as ctx: + job_resource = DataDesignerResource(ctx.sdk).get_job_resource(_JOB_NAME, workspace="default") + + with _patch_status(job_resource, "active"), caplog.at_level("INFO"): + analysis = job_resource.load_analysis() + + assert analysis.num_records == 3 + assert any("still cooking" in record.message.lower() for record in caplog.records) + + +@pytest.mark.asyncio +async def test_load_analysis_when_terminally_incomplete_warns_and_returns_partial( + caplog: pytest.LogCaptureFixture, +) -> None: + async with _completed_job() as ctx: + job_resource = DataDesignerResource(ctx.sdk).get_job_resource(_JOB_NAME, workspace="default") + + with _patch_status(job_resource, "error"), caplog.at_level("WARNING"): + analysis = job_resource.load_analysis() + + assert analysis.num_records == 3 + assert any("error" in record.message and "analysis" in record.message for record in caplog.records) + + +@pytest.mark.asyncio +async def test_load_analysis_raises_friendly_error_when_active_but_result_missing() -> None: + """An ``active`` job whose analysis result hasn't been written yet returns a 404 from the + Jobs service; ``_check_if_result_available`` translates that into a friendly + ``"'analysis' result is not available."`` message instead of leaking the underlying + HTTP error. + """ + async with _pending_job() as ctx: + job_resource = DataDesignerResource(ctx.sdk).get_job_resource(_JOB_NAME, workspace="default") + + with _patch_status(job_resource, "active"): + with pytest.raises(DataDesignerJobError, match="'analysis' result is not available"): + job_resource.load_analysis() diff --git a/plugins/nemo-data-designer/tests/integration/test_task.py b/plugins/nemo-data-designer/tests/integration/test_job_task.py similarity index 66% rename from plugins/nemo-data-designer/tests/integration/test_task.py rename to plugins/nemo-data-designer/tests/integration/test_job_task.py index af017ac5a5..68a6f6fee6 100644 --- a/plugins/nemo-data-designer/tests/integration/test_task.py +++ b/plugins/nemo-data-designer/tests/integration/test_job_task.py @@ -1,12 +1,27 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -import io +"""Integration tests for the Data Designer job *task* — the in-container worker +that the Jobs service invokes. + +These tests exercise ``nemo_data_designer_plugin.jobs.run`` end-to-end via +``task_context``: they run the task in-process, then verify the produced +results by reading them back through the high-level SDK +(``DataDesignerJobResource`` + ``DataDesignerJobResults``), which is the +recommended way for users to consume job artifacts. + +A note on patching ``get_job_status``: ``task_context`` invokes the task +directly and does not drive the Jobs service controller that would normally +roll the task's terminal state up to a ``completed`` job status. The SDK's +``download_artifacts`` gates on that status, so we patch it to ``"completed"`` +on the resource as the (single) network seam after the task has actually +written its results. +""" + import logging -import tarfile -import tempfile from collections.abc import Generator, Iterator from contextlib import contextmanager +from pathlib import Path from unittest.mock import patch import data_designer.config as dd @@ -17,19 +32,31 @@ from nemo_data_designer_plugin.jobs.run import BUFFER_SIZE from nemo_data_designer_plugin.jobs.spec import DataDesignerJobConfig from nemo_data_designer_plugin.jobs.task_results import ANALYSIS_RESULT_NAME, ARTIFACTS_RESULT_NAME +from nemo_data_designer_plugin.sdk.job_results import DataDesignerJobResults +from nemo_data_designer_plugin.sdk.resources import DataDesignerResource + +pytestmark = pytest.mark.integration + + +def _load_results(ctx: u.CreateJobTestContext, job_name: str, tmp_path: Path) -> DataDesignerJobResults: + """Read back a completed task's results through the high-level SDK. + ``task_context`` does not drive the Jobs service controller, so the platform + job status never rolls up to ``"completed"`` on its own. We patch the + resource's ``get_job_status`` to bypass that gate, then download artifacts + the same way an end user would. + """ + job_resource = DataDesignerResource(ctx.sdk).get_job_resource(job_name, workspace="default") + with patch.object(job_resource, "get_job_status", return_value="completed"): + return job_resource.download_artifacts(tmp_path) -def _get_dataset(ctx: u.CreateJobTestContext, job_name: str) -> pd.DataFrame: - with tempfile.TemporaryDirectory() as tmpdir: - artifacts_download = ctx.sdk.jobs.results.download("artifacts", job=job_name) - with tarfile.open(fileobj=io.BytesIO(artifacts_download.read()), mode="r:*") as tar: - tar.extractall(path=tmpdir) - return pd.read_parquet(f"{tmpdir}/artifacts/dataset/parquet-files") +def _get_dataset(ctx: u.CreateJobTestContext, job_name: str, tmp_path: Path) -> pd.DataFrame: + return _load_results(ctx, job_name, tmp_path).load_dataset() -def _get_analysis(ctx: u.CreateJobTestContext, job_name: str) -> DatasetProfilerResults: - response = ctx.sdk.jobs.results.download("analysis", job=job_name) - return DatasetProfilerResults.model_validate_json(response.read().decode()) + +def _get_analysis(ctx: u.CreateJobTestContext, job_name: str, tmp_path: Path) -> DatasetProfilerResults: + return _load_results(ctx, job_name, tmp_path).load_analysis() @pytest.fixture @@ -38,9 +65,8 @@ def _failing_result_manager() -> Generator[None]: yield -@pytest.mark.integration @pytest.mark.asyncio -async def test_task() -> None: +async def test_task(tmp_path: Path) -> None: test_value = "test-value" num_records = 42 builder = dd.DataDesignerConfigBuilder(model_configs=[u.make_model_config()]) @@ -64,20 +90,21 @@ async def test_task() -> None: assert ANALYSIS_RESULT_NAME in result_names assert ARTIFACTS_RESULT_NAME in result_names - dataset = _get_dataset(ctx, job_name) + dataset = _get_dataset(ctx, job_name, tmp_path) expected_partial_data = pd.DataFrame(data={"foo": [test_value] * num_records}) - pd.testing.assert_frame_equal(dataset, expected_partial_data) + # ``check_dtype=False``: the SDK loads via ``read_parquet_dataset``, which yields + # ``string[pyarrow]`` columns; the inline expected DataFrame is plain ``object``. + pd.testing.assert_frame_equal(dataset, expected_partial_data, check_dtype=False) - analysis = _get_analysis(ctx, job_name) + analysis = _get_analysis(ctx, job_name, tmp_path) assert analysis.num_records == 42 -@pytest.mark.integration @pytest.mark.asyncio @pytest.mark.skip( reason="Batch-completion artifact saves are not yet available through the high-level library interfaces" ) -async def test_save_partial_dataset_on_failure(_failing_result_manager: None) -> None: +async def test_save_partial_dataset_on_failure(_failing_result_manager: None, tmp_path: Path) -> None: test_value = "test-value" builder = dd.DataDesignerConfigBuilder(model_configs=[u.make_model_config()]) builder.add_column( @@ -104,10 +131,12 @@ async def test_save_partial_dataset_on_failure(_failing_result_manager: None) -> assert ANALYSIS_RESULT_NAME not in result_names assert ARTIFACTS_RESULT_NAME in result_names - dataset = _get_dataset(ctx, job_name) - pd.testing.assert_frame_equal(dataset, expected_partial_data) + dataset = _get_dataset(ctx, job_name, tmp_path) + pd.testing.assert_frame_equal(dataset, expected_partial_data, check_dtype=False) +# TODO: once we restore batch-completion artifact saves, we can drop this test +# and include the log-related assertion in that test instead (immediately above) @pytest.mark.asyncio async def test_exiting_with_error() -> None: builder = dd.DataDesignerConfigBuilder(model_configs=[u.make_model_config()]) @@ -131,20 +160,20 @@ async def test_exiting_with_error() -> None: @pytest.mark.asyncio -async def test_seed_dataset() -> None: +async def test_seed_dataset(tmp_path: Path) -> None: builder = dd.DataDesignerConfigBuilder(model_configs=[u.make_model_config()]) builder.with_seed_dataset(dd.HuggingFaceSeedSource(path="path/to/data.parquet")) builder.add_column(column_config=dd.ExpressionColumnConfig(name="full_name", expr=u.FULL_NAME_EXPR)) dd_job_config = DataDesignerJobConfig(num_records=3, config=builder.build()) - job_config = await u.compile_create_job(dd_job_config) + job_config = await u.compile_create_job(dd_job_config, workspace="default") job_name = "data-designer-abc123" with u.mock_hf_seed_reader(): async with u.task_context(job_config, job_name) as ctx: result = ctx.run_task() assert result.exit_code == 0 - dataset = _get_dataset(ctx, job_name) + dataset = _get_dataset(ctx, job_name, tmp_path) assert set(dataset["full_name"].values) == u.FULL_NAMES diff --git a/plugins/nemo-data-designer/tests/unit/test_cli.py b/plugins/nemo-data-designer/tests/integration/test_personas_cli.py similarity index 66% rename from plugins/nemo-data-designer/tests/unit/test_cli.py rename to plugins/nemo-data-designer/tests/integration/test_personas_cli.py index 4bb6eb73c1..8331f9242a 100644 --- a/plugins/nemo-data-designer/tests/unit/test_cli.py +++ b/plugins/nemo-data-designer/tests/integration/test_personas_cli.py @@ -4,20 +4,14 @@ from collections.abc import Generator from unittest.mock import Mock, patch +import nemo_data_designer_plugin.testing.utils as u import pytest -import typer from data_designer_nemo.nemotron_personas import WORKSPACE, get_resource_name_for_locale from nemo_data_designer_plugin.cli import personas as personas_module -from nemo_data_designer_plugin.cli.main import DataDesignerCLI -from nemo_data_designer_plugin.functions.preview import PreviewFunction -from nemo_data_designer_plugin.jobs.create import CreateJob from nemo_platform import NeMoPlatform from nemo_platform.types.files import NGCStorageConfig -from nemo_platform_plugin.commands import add_function_commands, add_job_commands -from nmp.core.files.service import FilesService -from nmp.core.secrets.service import SecretsService -from nmp.testing import create_test_client -from typer.testing import CliRunner + +pytestmark = pytest.mark.integration @pytest.fixture @@ -40,41 +34,28 @@ def mock_ngc_client() -> Generator[dict[str, Mock]]: @pytest.fixture def sdk(monkeypatch: pytest.MonkeyPatch, mock_ngc_client: dict[str, Mock]) -> Generator[NeMoPlatform]: - with create_test_client( - FilesService, - SecretsService, - client_type=NeMoPlatform, - ) as sdk: + with u.make_mock_client_context() as client_context: monkeypatch.setenv("NGC_API_KEY", "nvapi-abc123") - yield sdk + yield client_context.sdk monkeypatch.delenv("NGC_API_KEY") -@pytest.fixture -def runner() -> CliRunner: - return CliRunner() - - -@pytest.fixture -def app() -> typer.Typer: - cli = DataDesignerCLI() - typer_app = cli.get_cli() - add_function_commands(typer_app, {"preview": PreviewFunction}, cli=cli) - add_job_commands(typer_app, {"create": CreateJob}, cli=cli) - return typer_app - - @pytest.fixture def cli_sdk(monkeypatch: pytest.MonkeyPatch, sdk: NeMoPlatform) -> NeMoPlatform: monkeypatch.setattr(personas_module, "NeMoPlatform", lambda: sdk) return sdk -def test_make_fileset_creates_requested_locale_with_existing_secret( - runner: CliRunner, app: typer.Typer, cli_sdk: NeMoPlatform -) -> None: - result = runner.invoke( - app, +def test_personas_download_is_wired_properly() -> None: + result = u.invoke_cli(["personas", "download", "--help"]) + + assert result.exit_code == 0, result.output + assert "nemo data-designer personas download --list" in result.output + assert "data-designer download personas" not in result.output + + +def test_make_fileset_creates_requested_locale_with_existing_secret(cli_sdk: NeMoPlatform) -> None: + result = u.invoke_cli( [ "personas", "make-fileset", @@ -82,7 +63,7 @@ def test_make_fileset_creates_requested_locale_with_existing_secret( "en_US", "--api-key-secret", "system/ngc-api-key", - ], + ] ) assert result.exit_code == 0, result.output @@ -95,12 +76,11 @@ def test_make_fileset_creates_requested_locale_with_existing_secret( def test_make_fileset_creates_secret_from_env_then_fileset( - monkeypatch: pytest.MonkeyPatch, runner: CliRunner, app: typer.Typer, cli_sdk: NeMoPlatform + monkeypatch: pytest.MonkeyPatch, cli_sdk: NeMoPlatform ) -> None: monkeypatch.setenv("MY_NGC_API_KEY", "nvapi-from-env") - result = runner.invoke( - app, + result = u.invoke_cli( [ "personas", "make-fileset", @@ -110,7 +90,7 @@ def test_make_fileset_creates_secret_from_env_then_fileset( "system/my-ngc-key", "--api-key-env-var", "MY_NGC_API_KEY", - ], + ] ) assert result.exit_code == 0, result.output @@ -122,9 +102,8 @@ def test_make_fileset_creates_secret_from_env_then_fileset( assert fileset.storage.api_key_secret == "system/my-ngc-key" -def test_make_fileset_missing_env_var_is_clear(runner: CliRunner, app: typer.Typer) -> None: - result = runner.invoke( - app, +def test_make_fileset_missing_env_var() -> None: + result = u.invoke_cli( [ "personas", "make-fileset", @@ -134,7 +113,7 @@ def test_make_fileset_missing_env_var_is_clear(runner: CliRunner, app: typer.Typ "system/my-ngc-key", "--api-key-env-var", "MISSING_NGC_API_KEY", - ], + ] ) assert result.exit_code != 0 @@ -142,9 +121,8 @@ def test_make_fileset_missing_env_var_is_clear(runner: CliRunner, app: typer.Typ assert "not set or is empty" in result.output -def test_make_fileset_unknown_locale_is_clear(runner: CliRunner, app: typer.Typer) -> None: - result = runner.invoke( - app, +def test_make_fileset_unknown_locale() -> None: + result = u.invoke_cli( [ "personas", "make-fileset", @@ -152,7 +130,7 @@ def test_make_fileset_unknown_locale_is_clear(runner: CliRunner, app: typer.Type "de_DE", "--api-key-secret", "system/ngc-api-key", - ], + ] ) assert result.exit_code != 0 @@ -160,9 +138,8 @@ def test_make_fileset_unknown_locale_is_clear(runner: CliRunner, app: typer.Type assert "de_DE" in result.output -def test_make_fileset_bare_secret_name_is_clear(runner: CliRunner, app: typer.Typer) -> None: - result = runner.invoke( - app, +def test_make_fileset_bare_secret_name() -> None: + result = u.invoke_cli( [ "personas", "make-fileset", @@ -170,7 +147,7 @@ def test_make_fileset_bare_secret_name_is_clear(runner: CliRunner, app: typer.Ty "en_US", "--api-key-secret", "ngc-api-key", - ], + ] ) assert result.exit_code != 0 @@ -178,13 +155,12 @@ def test_make_fileset_bare_secret_name_is_clear(runner: CliRunner, app: typer.Ty def test_make_fileset_create_secret_conflict_does_not_create_fileset( - monkeypatch: pytest.MonkeyPatch, runner: CliRunner, app: typer.Typer, cli_sdk: NeMoPlatform + monkeypatch: pytest.MonkeyPatch, cli_sdk: NeMoPlatform ) -> None: cli_sdk.secrets.create(workspace="system", name="my-ngc-key", value="nvapi-existing") monkeypatch.setenv("MY_NGC_API_KEY", "nvapi-from-env") - result = runner.invoke( - app, + result = u.invoke_cli( [ "personas", "make-fileset", @@ -194,7 +170,7 @@ def test_make_fileset_create_secret_conflict_does_not_create_fileset( "system/my-ngc-key", "--api-key-env-var", "MY_NGC_API_KEY", - ], + ] ) assert result.exit_code == 1 @@ -203,20 +179,30 @@ def test_make_fileset_create_secret_conflict_does_not_create_fileset( assert filesets.data == [] -def test_nemotron_personas_download_is_wired(runner: CliRunner, app: typer.Typer) -> None: - result = runner.invoke(app, ["personas", "download", "--help"]) - - assert result.exit_code == 0, result.output - assert "Download Nemotron-Personas" in result.output - assert "nemo data-designer personas download --list" in result.output - assert "data-designer download personas" not in result.output - +def test_make_fileset_create_secret_internal_error_surfaces_clearly( + monkeypatch: pytest.MonkeyPatch, cli_sdk: NeMoPlatform +) -> None: + monkeypatch.setenv("MY_NGC_API_KEY", "nvapi-from-env") -@pytest.mark.parametrize("verb", ["run", "submit"]) -def test_preview_exposes_save_results_flags(runner: CliRunner, app: typer.Typer, verb: str) -> None: - result = runner.invoke(app, ["preview", verb, "--help"]) + def _boom(*args: object, **kwargs: object) -> None: + raise RuntimeError("secrets backend exploded") + + with patch.object(cli_sdk.secrets, "create", side_effect=_boom): + result = u.invoke_cli( + [ + "personas", + "make-fileset", + "--locale", + "en_US", + "--api-key-secret", + "system/my-ngc-key", + "--api-key-env-var", + "MY_NGC_API_KEY", + ] + ) - assert result.exit_code == 0, result.output - assert "--save-results" in result.output - assert "--artifact-path" in result.output - assert "--non-interactive" in result.output + assert result.exit_code == 1 + assert "Failed to create secret" in result.output + assert "secrets backend exploded" in result.output + filesets = cli_sdk.files.filesets.list(workspace=WORKSPACE) + assert filesets.data == [] diff --git a/plugins/nemo-data-designer/tests/integration/test_cli_local.py b/plugins/nemo-data-designer/tests/integration/test_preview_local_cli.py similarity index 80% rename from plugins/nemo-data-designer/tests/integration/test_cli_local.py rename to plugins/nemo-data-designer/tests/integration/test_preview_local_cli.py index 326fe1b85b..ac4cc78d75 100644 --- a/plugins/nemo-data-designer/tests/integration/test_cli_local.py +++ b/plugins/nemo-data-designer/tests/integration/test_preview_local_cli.py @@ -6,30 +6,18 @@ import nemo_data_designer_plugin.testing.utils as u import pandas as pd import pytest -import typer from data_designer.cli.utils.sample_records_pager import PAGER_FILENAME from data_designer.config.analysis.dataset_profiler import DatasetProfilerResults -from typer.testing import CliRunner +pytestmark = pytest.mark.integration -@pytest.fixture -def runner() -> CliRunner: - return CliRunner() - -@pytest.fixture -def app() -> typer.Typer: - return u.make_data_designer_cli_app() - - -@pytest.mark.integration -def test_preview_run_saves_expected_artifacts(runner: CliRunner, app: typer.Typer, tmp_path: Path) -> None: +def test_preview_run_saves_expected_artifacts(tmp_path: Path) -> None: config_path = _write_sampler_config(tmp_path) artifact_path = tmp_path / "preview-artifacts" with u.make_mock_client_context(workspace="default") as client_context: - result = runner.invoke( - app, + result = u.invoke_cli( [ "preview", "run", @@ -40,7 +28,7 @@ def test_preview_run_saves_expected_artifacts(runner: CliRunner, app: typer.Type "--artifact-path", str(artifact_path), ], - obj=u.make_data_designer_cli_state(client_context), + client_context, ) assert result.exit_code == 0, result.output @@ -53,8 +41,7 @@ def test_preview_run_saves_expected_artifacts(runner: CliRunner, app: typer.Type assert (results_dir / "sample_records" / PAGER_FILENAME).exists() -@pytest.mark.integration -def test_preview_run_supports_local_file_seed_source(runner: CliRunner, app: typer.Typer, tmp_path: Path) -> None: +def test_preview_run_supports_local_file_seed_source(tmp_path: Path) -> None: seed_path = tmp_path / "seed.parquet" u.SEED_DATA.to_parquet(seed_path, index=False) config_path = u.write_config_file( @@ -74,8 +61,7 @@ def load_config_builder() -> dd.DataDesignerConfigBuilder: artifact_path = tmp_path / "preview-artifacts" with u.make_mock_client_context(workspace="default") as client_context: - result = runner.invoke( - app, + result = u.invoke_cli( [ "preview", "run", @@ -86,7 +72,7 @@ def load_config_builder() -> dd.DataDesignerConfigBuilder: "--artifact-path", str(artifact_path), ], - obj=u.make_data_designer_cli_state(client_context), + client_context, ) assert result.exit_code == 0, result.output @@ -94,10 +80,7 @@ def load_config_builder() -> dd.DataDesignerConfigBuilder: assert set(dataset["full_name"].tolist()) == u.FULL_NAMES -@pytest.mark.integration -def test_preview_run_rejects_dataframe_seed_with_clear_error( - runner: CliRunner, app: typer.Typer, tmp_path: Path -) -> None: +def test_preview_run_rejects_dataframe_seed_with_clear_error(tmp_path: Path) -> None: config_path = u.write_config_file( tmp_path, """ @@ -115,10 +98,9 @@ def load_config_builder() -> dd.DataDesignerConfigBuilder: ) with u.make_mock_client_context(workspace="default") as client_context: - result = runner.invoke( - app, + result = u.invoke_cli( ["preview", "run", str(config_path), "--num-records", "3"], - obj=u.make_data_designer_cli_state(client_context), + client_context, ) message = result.output @@ -130,15 +112,14 @@ def load_config_builder() -> dd.DataDesignerConfigBuilder: assert "No such file" not in message -@pytest.mark.integration -def test_create_run_reports_artifacts_and_dataset_path(runner: CliRunner, app: typer.Typer, tmp_path: Path) -> None: +def test_create_run_reports_artifacts_and_dataset_path(tmp_path: Path) -> None: config_path = _write_sampler_config(tmp_path) with u.make_mock_client_context(workspace="default") as client_context: - result = runner.invoke( - app, + result = u.invoke_cli( ["create", "run", str(config_path), "--num-records", "3"], - obj=u.make_data_designer_cli_state(client_context, output_format="json"), + client_context, + output_format="json", ) assert result.exit_code == 0, result.output diff --git a/plugins/nemo-data-designer/tests/integration/test_preview.py b/plugins/nemo-data-designer/tests/integration/test_preview_remote_sdk.py similarity index 85% rename from plugins/nemo-data-designer/tests/integration/test_preview.py rename to plugins/nemo-data-designer/tests/integration/test_preview_remote_sdk.py index 48e0ad8842..c1d7fbaef5 100644 --- a/plugins/nemo-data-designer/tests/integration/test_preview.py +++ b/plugins/nemo-data-designer/tests/integration/test_preview_remote_sdk.py @@ -11,10 +11,11 @@ import pytest from data_designer_nemo.fileset_file_seed_source import FilesetFileSeedSource from nemo_data_designer_plugin.config import get_config -from nemo_data_designer_plugin.sdk.errors import DataDesignerConfigValidationError +from nemo_data_designer_plugin.sdk.errors import DataDesignerConfigValidationError, DataDesignerPreviewError + +pytestmark = pytest.mark.integration -@pytest.mark.integration def test_request_too_many_records() -> None: too_many_records = get_config().preview_num_records.max + 1 @@ -36,7 +37,6 @@ def test_request_too_many_records() -> None: assert "Max num records" in str(exc_info.value) -@pytest.mark.integration def test_happy_path_preview() -> None: column_name = "column-name" value = "a" @@ -67,8 +67,7 @@ def test_happy_path_preview() -> None: assert_message_with(log_messages, fuzzy="Preview generation in progress") -@pytest.mark.integration -def test_seed_dataset() -> None: +def test_hf_seed_dataset() -> None: builder = dd.DataDesignerConfigBuilder(model_configs=[u.make_model_config()]) builder.with_seed_dataset( dd.HuggingFaceSeedSource(path="my-ws/my-fileset#path/to/data.parquet", token=u.SECRET_NAME) @@ -87,7 +86,6 @@ def test_seed_dataset() -> None: assert set(preview_results.dataset["full_name"].values) == u.FULL_NAMES -@pytest.mark.integration def test_fileset_file_seed_dataset_plugin() -> None: builder = dd.DataDesignerConfigBuilder(model_configs=[u.make_model_config()]) builder.with_seed_dataset(FilesetFileSeedSource(path=u.FILESET_FILE_SEED_SOURCE_PATH)) # ty: ignore[invalid-argument-type] @@ -104,7 +102,6 @@ def test_fileset_file_seed_dataset_plugin() -> None: assert set(preview_results.dataset["full_name"].values) == u.FULL_NAMES -@pytest.mark.integration def test_nemotron_personas_dataset() -> None: builder = dd.DataDesignerConfigBuilder(model_configs=[u.make_model_config()]) builder.add_column( @@ -137,7 +134,6 @@ def _parse_age(v: str) -> int: assert all(25 <= age <= 45 for age in demo_ages) -@pytest.mark.integration def test_preview_with_schema_transform_processor() -> None: column_name = "school_subject" processor_name = "chat_format" @@ -175,6 +171,33 @@ def test_preview_with_schema_transform_processor() -> None: assert "messages" in processor_records[0] +def test_preview_surfaces_worker_error_through_sdk(monkeypatch: pytest.MonkeyPatch) -> None: + """When the preview worker thread raises, the function emits a ``LogFrame`` and an + ``Error`` frame instead of ``Done``; the SDK's ``_PreviewFrameCollector`` translates + that ``Error`` into a typed ``DataDesignerPreviewError`` with the original message. + """ + from nemo_data_designer_plugin.functions import _preview_worker as worker_module + + def boom(*args: object, **kwargs: object) -> None: + raise RuntimeError("forced worker failure") + + monkeypatch.setattr(worker_module, "make_preview_dataset", boom) + + builder = dd.DataDesignerConfigBuilder(model_configs=[u.make_model_config()]) + builder.add_column( + column_config=dd.SamplerColumnConfig( + name="foo", + sampler_type=dd.SamplerType.CATEGORY, + params=dd.CategorySamplerParams(values=["a"]), + ) + ) + + with u.make_mock_client_context() as client_context: + dd_client = u.make_dd_client(client_context) + with pytest.raises(DataDesignerPreviewError, match="forced worker failure"): + dd_client.preview(builder, num_records=3) + + def assert_message_with(messages: list[str], exact: str | None = None, fuzzy: str | None = None) -> None: match (exact, fuzzy): case (None, None): diff --git a/plugins/nemo-data-designer/tests/integration/test_validation.py b/plugins/nemo-data-designer/tests/integration/test_remote_validation_errors.py similarity index 82% rename from plugins/nemo-data-designer/tests/integration/test_validation.py rename to plugins/nemo-data-designer/tests/integration/test_remote_validation_errors.py index 56298e2f40..1b6cd1a059 100644 --- a/plugins/nemo-data-designer/tests/integration/test_validation.py +++ b/plugins/nemo-data-designer/tests/integration/test_remote_validation_errors.py @@ -7,6 +7,8 @@ import pytest from nemo_data_designer_plugin.sdk.errors import DataDesignerClientError, DataDesignerConfigValidationError +pytestmark = pytest.mark.integration + def _assert_error( dd_client, @@ -25,12 +27,8 @@ def _assert_error( assert fragment in str(exc_info.value) -@pytest.mark.integration -def test_unknown_provider_in_request() -> None: - unknown_provider = "some-unknown-provider" - bad_model_config = u.make_model_config(provider=unknown_provider) - - builder = dd.DataDesignerConfigBuilder(model_configs=[bad_model_config]) +def _builder_with_llm_column(model_config: dd.ModelConfig) -> dd.DataDesignerConfigBuilder: + builder = dd.DataDesignerConfigBuilder(model_configs=[model_config]) builder.add_column( column_config=dd.SamplerColumnConfig( name="foo", sampler_type=dd.SamplerType.CATEGORY, params=dd.CategorySamplerParams(values=["a", "b"]) @@ -38,31 +36,47 @@ def test_unknown_provider_in_request() -> None: ) builder.add_column( column_config=dd.LLMTextColumnConfig( - name="story", prompt="Write a story about {{ foo }}", model_alias=bad_model_config.alias + name="story", prompt="Write a story about {{ foo }}", model_alias=model_config.alias ) ) + return builder + + +def test_unknown_provider_in_request() -> None: + unknown_provider = "some-unknown-provider" + bad_model_config = u.make_model_config(provider=unknown_provider) + builder = _builder_with_llm_column(bad_model_config) with u.make_mock_client_context() as client_context: dd_client = u.make_dd_client(client_context) _assert_error(dd_client, builder, ["Cannot access provider", unknown_provider]) -@pytest.mark.integration +def test_model_config_without_explicit_provider_is_rejected() -> None: + alias = "no-provider-specified" + bad_model_config = dd.ModelConfig(alias=alias, model="some-model") + builder = _builder_with_llm_column(bad_model_config) + + with u.make_mock_client_context() as client_context: + dd_client = u.make_dd_client(client_context) + _assert_error(dd_client, builder, ["does not have an explicit provider defined", alias]) + + +def test_malformed_provider_reference_is_rejected() -> None: + alias = "too-many-slashes" + malformed_provider_name = "foo/bar/baz" + bad_model_config = dd.ModelConfig(alias=alias, model="some-model", provider=malformed_provider_name) + builder = _builder_with_llm_column(bad_model_config) + + with u.make_mock_client_context() as client_context: + dd_client = u.make_dd_client(client_context) + _assert_error(dd_client, builder, ["Malformed model provider", alias, malformed_provider_name]) + + def test_invalid_models_provided() -> None: disallowed_model = "this-model-is-not-allowed" bad_model_config = u.make_model_config(provider=u.RESTRICTED_PROVIDER_NAME, model=disallowed_model) - - builder = dd.DataDesignerConfigBuilder(model_configs=[bad_model_config]) - builder.add_column( - column_config=dd.SamplerColumnConfig( - name="foo", sampler_type=dd.SamplerType.CATEGORY, params=dd.CategorySamplerParams(values=["a", "b"]) - ) - ) - builder.add_column( - column_config=dd.LLMTextColumnConfig( - name="story", prompt="Write a story about {{ foo }}", model_alias=bad_model_config.alias - ) - ) + builder = _builder_with_llm_column(bad_model_config) with ( u.make_mock_client_context() as client_context, @@ -72,7 +86,6 @@ def test_invalid_models_provided() -> None: _assert_error(dd_client, builder, [disallowed_model, "not enabled for provider", u.RESTRICTED_PROVIDER_NAME]) -@pytest.mark.integration def test_unrecognized_model_alias() -> None: model_alias = "unknown-model-alias" @@ -97,7 +110,6 @@ def test_unrecognized_model_alias() -> None: _assert_error(dd_client, builder, ["Unrecognized", model_alias]) -@pytest.mark.integration def test_mcp_tools_not_allowed() -> None: provider = u.OPEN_PROVIDER_NAME model_config = u.make_model_config(provider=provider) @@ -116,7 +128,6 @@ def test_mcp_tools_not_allowed() -> None: _assert_error(dd_client, builder, ["Tool configs are not supported"]) -@pytest.mark.integration def test_seed_dataset_bad_token() -> None: bad_token_secret = "unrecognized-secret-ref" builder = dd.DataDesignerConfigBuilder(model_configs=[u.make_model_config()]) @@ -133,7 +144,6 @@ def test_seed_dataset_bad_token() -> None: _assert_error(dd_client, builder, [bad_token_secret]) -@pytest.mark.integration def test_nemotron_personas_dataset_failure() -> None: builder = dd.DataDesignerConfigBuilder(model_configs=[u.make_model_config()]) builder.add_column( @@ -150,7 +160,6 @@ def test_nemotron_personas_dataset_failure() -> None: _assert_error(dd_client, builder, ["Nemotron personas filesets"], DataDesignerClientError) -@pytest.mark.integration def test_server_side_unsupported_seed_type_validation() -> None: builder = dd.DataDesignerConfigBuilder(model_configs=[u.make_model_config()]) builder.with_seed_dataset(dd.DataFrameSeedSource(df=pd.DataFrame(data={"a": [1, 2, 3]}))) diff --git a/plugins/nemo-data-designer/tests/integration/test_sdk_get_model_providers.py b/plugins/nemo-data-designer/tests/integration/test_sdk_get_model_providers.py new file mode 100644 index 0000000000..703917888a --- /dev/null +++ b/plugins/nemo-data-designer/tests/integration/test_sdk_get_model_providers.py @@ -0,0 +1,36 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import data_designer.config as dd +import nemo_data_designer_plugin.testing.utils as u +import pytest + +pytestmark = pytest.mark.integration + + +def test_get_default_model_providers_returns_registered_providers() -> None: + """The SDK exposes IGW-registered providers as Data Designer ModelProviders.""" + + with ( + u.make_mock_client_context() as client_context, + u.setup_mock_providers(client_context), + ): + dd_client = u.make_dd_client(client_context) + providers = dd_client.get_default_model_providers() + + provider_names = {provider.name for provider in providers} + assert u.OPEN_PROVIDER_NAME in {name.split("/")[-1] for name in provider_names} + assert u.RESTRICTED_PROVIDER_NAME in {name.split("/")[-1] for name in provider_names} + for provider in providers: + assert isinstance(provider, dd.ModelProvider) + assert provider.endpoint, f"Provider {provider.name!r} has no endpoint" + + +def test_get_default_model_providers_returns_empty_list_when_none_registered() -> None: + """No registered providers means the SDK returns an empty list (not None, not an error).""" + + with u.make_mock_client_context() as client_context: + dd_client = u.make_dd_client(client_context) + providers = dd_client.get_default_model_providers() + + assert providers == [] diff --git a/plugins/nemo-data-designer/tests/unit/test_errors.py b/plugins/nemo-data-designer/tests/unit/test_errors.py new file mode 100644 index 0000000000..2b6a4de0df --- /dev/null +++ b/plugins/nemo-data-designer/tests/unit/test_errors.py @@ -0,0 +1,87 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for the small error-extraction helper that replaces the legacy +``if "404" in str(e)`` pattern. Because the helper is pure logic over an +``httpx.Response``, it sits naturally as a unit test.""" + +import httpx +import pytest +from nemo_data_designer_plugin.sdk.errors import ( + DataDesignerClientError, + DataDesignerJobError, + extract_http_error_info, +) + + +def _make_status_error(status_code: int, *, body: str | None = None, json_body: object = None) -> httpx.HTTPStatusError: + request = httpx.Request("GET", "http://testserver/whatever") + if json_body is not None: + response = httpx.Response(status_code, request=request, json=json_body) + else: + response = httpx.Response(status_code, request=request, text=body or "") + return httpx.HTTPStatusError("err", request=request, response=response) + + +def test_extract_http_error_info_pulls_detail_field_from_json_body() -> None: + exc = _make_status_error(404, json_body={"detail": "Job result not found"}) + + status_code, detail = extract_http_error_info(exc) + + assert status_code == 404 + assert detail == "Job result not found" + + +def test_extract_http_error_info_falls_back_to_raw_body_when_not_json() -> None: + exc = _make_status_error(500, body="oh no") + + status_code, detail = extract_http_error_info(exc) + + assert status_code == 500 + assert detail == "oh no" + + +def test_extract_http_error_info_falls_back_to_raw_body_when_json_lacks_detail() -> None: + exc = _make_status_error(422, json_body={"errors": ["bad config"]}) + + status_code, detail = extract_http_error_info(exc) + + assert status_code == 422 + # Body is JSON but has no string ``detail`` field, so we surface the raw text. + assert "errors" in detail + + +def test_extract_http_error_info_ignores_non_string_detail_field() -> None: + """A ``detail`` field that isn't a string (e.g. a list of validation errors) should + fall through to the raw body so callers always get a string.""" + exc = _make_status_error(422, json_body={"detail": [{"loc": ["body"], "msg": "field required"}]}) + + _, detail = extract_http_error_info(exc) + + assert isinstance(detail, str) + assert "field required" in detail + + +def test_data_designer_client_error_carries_status_code() -> None: + err = DataDesignerJobError("Job result not found", status_code=404) + assert isinstance(err, DataDesignerClientError) + assert err.status_code == 404 + assert str(err) == "Job result not found" + + +def test_data_designer_client_error_status_code_defaults_to_none() -> None: + """Locally-constructed errors (where there's no upstream HTTP response) carry no status.""" + err = DataDesignerJobError("Current job status is 'cancelled', results are not available.") + assert err.status_code is None + + +@pytest.mark.parametrize("status_code", [404, 422, 500]) +def test_status_code_round_trips_through_raise_from(status_code: int) -> None: + """Make sure the status code survives the ``raise X from exc`` pattern used in + ``_raise_for_status`` / ``_get_error``.""" + cause = _make_status_error(status_code, json_body={"detail": "boom"}) + try: + raise DataDesignerJobError("boom", status_code=status_code) from cause + except DataDesignerJobError as e: + assert e.status_code == status_code + assert e.__cause__ is cause diff --git a/plugins/nemo-data-designer/tests/unit/test_job_results.py b/plugins/nemo-data-designer/tests/unit/test_job_results.py deleted file mode 100644 index 63d1a27a01..0000000000 --- a/plugins/nemo-data-designer/tests/unit/test_job_results.py +++ /dev/null @@ -1,57 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -from pathlib import Path -from unittest.mock import Mock - -import pandas as pd -import pytest -from nemo_data_designer_plugin.sdk.errors import ( - DataDesignerClientError, - DataDesignerJobError, -) -from nemo_data_designer_plugin.sdk.job_results import DataDesignerJobResults - - -def test_load_analysis_success(tmp_path: Path) -> None: - mock_analysis = Mock() - results = DataDesignerJobResults(artifacts_dir=tmp_path, analysis=mock_analysis) - assert results.load_analysis() is mock_analysis - - -def test_load_analysis_raises_when_error_string(tmp_path: Path) -> None: - error_msg = "Unable to fetch analysis: something went wrong" - results = DataDesignerJobResults(artifacts_dir=tmp_path, analysis=error_msg) - with pytest.raises(DataDesignerJobError, match=error_msg): - results.load_analysis() - - -def test_load_dataset(tmp_path: Path) -> None: - parquet_files_dir = tmp_path / "dataset" / "parquet-files" - parquet_files_dir.mkdir(parents=True) - - expected_df = pd.DataFrame({"col": [1, 2, 3]}).convert_dtypes(dtype_backend="pyarrow") - expected_df.to_parquet(f"{parquet_files_dir}/00000.parquet", index=False) - - results = DataDesignerJobResults(artifacts_dir=tmp_path, analysis=Mock()) - - dataset = results.load_dataset() - pd.testing.assert_frame_equal(dataset, expected_df) - - -def test_load_processor_dataset(tmp_path: Path) -> None: - processor_name = "chat_format" - - processor_files_dir = tmp_path / "dataset" / "processors-files" / processor_name - processor_files_dir.mkdir(parents=True) - - expected_df = pd.DataFrame({"col": [1, 2, 3]}).convert_dtypes(dtype_backend="pyarrow") - expected_df.to_parquet(f"{processor_files_dir}/00000.parquet", index=False) - - results = DataDesignerJobResults(artifacts_dir=tmp_path, analysis=Mock()) - - dataset = results.load_processor_dataset(processor_name) - pd.testing.assert_frame_equal(dataset, expected_df) - - with pytest.raises(DataDesignerClientError): - results.load_processor_dataset("undefined-processor") diff --git a/plugins/nemo-data-designer/tests/unit/test_model_provider.py b/plugins/nemo-data-designer/tests/unit/test_model_provider.py index 9f5f0bb279..4a0b8449eb 100644 --- a/plugins/nemo-data-designer/tests/unit/test_model_provider.py +++ b/plugins/nemo-data-designer/tests/unit/test_model_provider.py @@ -14,35 +14,12 @@ ) -@pytest.mark.asyncio -async def test_provider_cannot_be_none() -> None: - alias = "no-provider-specified" - bad_model_configs = [ - dd.ModelConfig( - alias=alias, - model="some-model", - ) - ] - - with ( - u.make_mock_client_context() as client_context, - pytest.raises(NDDInvalidConfigError) as exc_info, - ): - await make_model_provider_registry( - bad_model_configs, sdk=client_context.async_sdk, default_workspace=u.WORKSPACE_NAME - ) - assert "does not have an explicit provider defined" in str(exc_info.value) - - @pytest.mark.asyncio async def test_local_first_provider_cannot_be_none() -> None: - alias = "no-provider-specified" - bad_model_configs = [ - dd.ModelConfig( - alias=alias, - model="some-model", - ) - ] + """When a local-first registry build sees a missing provider, it must fail fast + *before* hitting the local-provider lookup helper. + """ + bad_model_configs = [dd.ModelConfig(alias="no-provider-specified", model="some-model")] with u.make_mock_client_context() as client_context: with ( @@ -57,118 +34,24 @@ async def test_local_first_provider_cannot_be_none() -> None: default_lookup.assert_not_called() assert "explicit provider defined" in str(exc_info.value) - assert "Missing provider(s): []" not in str(exc_info.value) @pytest.mark.asyncio -async def test_malformed_provider_name() -> None: - alias = "too-many-slashes" - malformed_provider_name = "foo/bar/baz" - bad_model_configs = [ - dd.ModelConfig( - alias=alias, - model="some-model", - provider=malformed_provider_name, - ) - ] - - with ( - u.make_mock_client_context() as client_context, - pytest.raises(NDDInvalidConfigError) as exc_info, - ): - await make_model_provider_registry( - bad_model_configs, sdk=client_context.async_sdk, default_workspace=u.WORKSPACE_NAME - ) - assert "Malformed model provider" in str(exc_info.value) - assert alias in str(exc_info.value) - assert malformed_provider_name in str(exc_info.value) - - -@pytest.mark.asyncio -async def test_inaccessible_provider() -> None: - inaccessible_provider_name = "inaccessible/provider" - bad_model_configs = [ - dd.ModelConfig( - alias="text", - model="some-model", - provider=inaccessible_provider_name, - ) - ] - - with ( - u.make_mock_client_context() as client_context, - pytest.raises(NDDInvalidConfigError) as exc_info, - ): - await make_model_provider_registry( - bad_model_configs, sdk=client_context.async_sdk, default_workspace=u.WORKSPACE_NAME - ) - assert "Cannot access provider" in str(exc_info.value) - assert inaccessible_provider_name in str(exc_info.value) - - -@pytest.mark.asyncio -async def test_disallowed_model_on_provider() -> None: - disallowed_model = "some-model-not-in-enabled-models-list" - model_configs = [ - dd.ModelConfig( - alias="text", - model=disallowed_model, - provider=u.RESTRICTED_PROVIDER_NAME, - ) - ] - - with ( - u.make_mock_client_context() as client_context, - u.setup_mock_providers(client_context), - pytest.raises(NDDInvalidConfigError) as exc_info, - ): - await make_model_provider_registry( - model_configs, sdk=client_context.async_sdk, default_workspace=u.WORKSPACE_NAME - ) - assert "not enabled for provider" in str(exc_info.value) - assert disallowed_model in str(exc_info.value) - - -@pytest.mark.asyncio -async def test_happy_path() -> None: - model_configs = [ - dd.ModelConfig( - alias="text", - model="anything", - provider=u.OPEN_PROVIDER_NAME, - ), - dd.ModelConfig( - alias="judge", - model=u.ENABLED_MODEL_NAME, - provider=u.RESTRICTED_PROVIDER_NAME, - ), - ] - - with ( - u.make_mock_client_context() as client_context, - u.setup_mock_providers(client_context), - ): - registry = await make_model_provider_registry( - model_configs, sdk=client_context.async_sdk, default_workspace=u.WORKSPACE_NAME - ) - - assert registry is not None - assert len(registry.providers) == 2 - expected_provider_names = {u.OPEN_PROVIDER_NAME, u.RESTRICTED_PROVIDER_NAME} - assert expected_provider_names == {provider.name for provider in registry.providers} - assert registry.default in expected_provider_names - - -@pytest.mark.asyncio -async def test_no_model_configs() -> None: +async def test_no_model_configs_returns_none() -> None: + """``make_model_provider_registry`` returns None for an empty model-config list, + which the engine treats as 'no LLMs in this config'. + """ with u.make_mock_client_context() as client_context: - assert ( - await make_model_provider_registry([], sdk=client_context.async_sdk, default_workspace=u.WORKSPACE_NAME) - is None + registry = await make_model_provider_registry( + [], sdk=client_context.async_sdk, default_workspace=u.WORKSPACE_NAME ) + assert registry is None def test_null_registry() -> None: + """Configs with no LLM columns get a one-provider 'no-op' registry so the + Data Designer engine can run without rejecting an empty registry. + """ registry = make_null_registry() assert len(registry.providers) == 1 diff --git a/plugins/nemo-data-designer/tests/unit/test_preview_function.py b/plugins/nemo-data-designer/tests/unit/test_preview_function.py index 6c9c888495..efc2f26dca 100644 --- a/plugins/nemo-data-designer/tests/unit/test_preview_function.py +++ b/plugins/nemo-data-designer/tests/unit/test_preview_function.py @@ -68,29 +68,6 @@ def fake_worker( assert [frame.model_dump()["kind"] for frame in frames] == ["log", "done"] -@pytest.mark.asyncio -async def test_preview_function_error_frame_terminates_without_done(monkeypatch: pytest.MonkeyPatch) -> None: - _patch_preview_dependencies(monkeypatch) - - def fake_worker(*args: object) -> None: - raise RuntimeError("boom") - - monkeypatch.setattr(worker_module, "make_preview_dataset", fake_worker) - - frames = [ - frame - async for frame in PreviewFunction().run( - PreviewSpec(config=_config(), num_records=2), - ctx=FunctionContext(workspace="team-a"), - async_sdk=AsyncMock(spec=AsyncNeMoPlatform), - is_local=True, - ) - ] - - assert [frame.model_dump()["kind"] for frame in frames] == ["log", "error"] - assert frames[-1].model_dump()["message"] == "boom" - - def test_preview_route_streams_ndjson_and_heartbeats(monkeypatch: pytest.MonkeyPatch) -> None: _patch_preview_dependencies(monkeypatch) diff --git a/plugins/nemo-data-designer/tests/unit/test_sdk_job_resources.py b/plugins/nemo-data-designer/tests/unit/test_sdk_job_resources.py deleted file mode 100644 index 19bbe3a2c3..0000000000 --- a/plugins/nemo-data-designer/tests/unit/test_sdk_job_resources.py +++ /dev/null @@ -1,159 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -import io -import json -import tarfile -from pathlib import Path -from unittest.mock import AsyncMock, patch - -import httpx -import pytest -import respx -from nemo_data_designer_plugin.sdk.errors import DataDesignerJobError -from nemo_data_designer_plugin.sdk.job_resources import AsyncDataDesignerJobResource, DataDesignerJobResource -from nemo_data_designer_plugin.sdk.job_results import DataDesignerJobResults -from nemo_platform import AsyncNeMoPlatform, NeMoPlatform - - -@pytest.fixture -def platform() -> NeMoPlatform: - return NeMoPlatform(base_url="http://testserver", workspace="test-workspace", access_token="token") - - -@pytest.fixture -def async_platform() -> AsyncNeMoPlatform: - return AsyncNeMoPlatform(base_url="http://testserver", workspace="test-workspace", access_token="token") - - -@pytest.fixture -def job_resource(platform: NeMoPlatform) -> DataDesignerJobResource: - return DataDesignerJobResource(job_name="test-job", platform=platform, workspace="test-workspace") - - -@pytest.fixture -def async_job_resource(async_platform: AsyncNeMoPlatform) -> AsyncDataDesignerJobResource: - return AsyncDataDesignerJobResource(job_name="test-job", platform=async_platform, workspace="test-workspace") - - -@respx.mock -def test_get_job(job_resource: DataDesignerJobResource) -> None: - respx.get("http://testserver/apis/data-designer/v2/workspaces/test-workspace/jobs/create/test-job").mock( - return_value=httpx.Response(200, json={"name": "test-job"}) - ) - assert job_resource.get_job()["name"] == "test-job" - - -@respx.mock -def test_get_job_status(job_resource: DataDesignerJobResource) -> None: - respx.get("http://testserver/apis/data-designer/v2/workspaces/test-workspace/jobs/create/test-job/status").mock( - return_value=httpx.Response(200, json={"status": "active"}) - ) - assert job_resource.get_job_status() == "active" - - -def test_check_if_complete_raises_when_not_complete(job_resource: DataDesignerJobResource) -> None: - with patch.object(job_resource, "get_job_status", return_value="active"): - with pytest.raises(DataDesignerJobError): - job_resource.check_if_complete(raise_if_not_complete=True) - - -def test_wait_until_done_success(job_resource: DataDesignerJobResource, caplog: pytest.LogCaptureFixture) -> None: - with ( - patch("nemo_data_designer_plugin.sdk.job_resources._pause"), - patch.object(job_resource, "get_job_status", side_effect=["active", "completed"]), - patch.object(job_resource, "get_logs", return_value=[]), - caplog.at_level("INFO"), - ): - job_resource.wait_until_done() - - assert any("completed successfully" in record.message for record in caplog.records) - - -@respx.mock -def test_get_logs_multiple_pages(job_resource: DataDesignerJobResource) -> None: - route = respx.get("http://testserver/apis/data-designer/v2/workspaces/test-workspace/jobs/create/test-job/logs") - route.mock( - side_effect=[ - httpx.Response( - 200, - json={ - "data": [ - { - "message": json.dumps( - {"name": "data_designer.something", "levelname": "INFO", "message": "Page 1"} - ) - } - ], - "next_page": "cursor1", - }, - ), - httpx.Response( - 200, - json={ - "data": [ - { - "message": json.dumps( - {"name": "data_designer.something", "levelname": "INFO", "message": "Page 2"} - ) - } - ], - "next_page": None, - }, - ), - ] - ) - - logs = job_resource.get_logs() - assert [log["message"] for log in logs] == ["Page 1", "Page 2"] - - -def _make_tar_bytes() -> bytes: - buffer = io.BytesIO() - with tarfile.open(fileobj=buffer, mode="w") as tar: - data = b"dummy" - info = tarfile.TarInfo(name="artifacts/dataset/parquet-files/00000.parquet") - info.size = len(data) - tar.addfile(info, io.BytesIO(data)) - return buffer.getvalue() - - -@respx.mock -def test_download_artifacts_success(job_resource: DataDesignerJobResource, tmp_path: Path) -> None: - with patch.object(job_resource, "_check_if_result_available"): - respx.get( - "http://testserver/apis/data-designer/v2/workspaces/test-workspace/jobs/create/test-job/results/artifacts/download" - ).mock(return_value=httpx.Response(200, content=_make_tar_bytes())) - respx.get( - "http://testserver/apis/data-designer/v2/workspaces/test-workspace/jobs/create/test-job/results/analysis/download" - ).mock( - return_value=httpx.Response(200, json={"num_records": 0, "target_num_records": 0, "column_statistics": []}) - ) - result = job_resource.download_artifacts(tmp_path) - - assert isinstance(result, DataDesignerJobResults) - - -@pytest.mark.asyncio -@respx.mock -async def test_get_job_async(async_job_resource: AsyncDataDesignerJobResource) -> None: - respx.get("http://testserver/apis/data-designer/v2/workspaces/test-workspace/jobs/create/test-job").mock( - return_value=httpx.Response(200, json={"name": "test-job"}) - ) - result = await async_job_resource.get_job() - assert result["name"] == "test-job" - - -@pytest.mark.asyncio -async def test_wait_until_done_success_async( - async_job_resource: AsyncDataDesignerJobResource, caplog: pytest.LogCaptureFixture -) -> None: - with ( - patch("nemo_data_designer_plugin.sdk.job_resources._async_pause"), - patch.object(async_job_resource, "get_job_status", new=AsyncMock(side_effect=["active", "completed"])), - patch.object(async_job_resource, "get_logs", new=AsyncMock(return_value=[])), - caplog.at_level("INFO"), - ): - await async_job_resource.wait_until_done() - - assert any("completed successfully" in record.message for record in caplog.records) diff --git a/plugins/nemo-data-designer/tests/unit/test_sdk_resources.py b/plugins/nemo-data-designer/tests/unit/test_sdk_resources.py index 956f4ace82..c07ea9b5ba 100644 --- a/plugins/nemo-data-designer/tests/unit/test_sdk_resources.py +++ b/plugins/nemo-data-designer/tests/unit/test_sdk_resources.py @@ -1,9 +1,21 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -import json +"""Pure-logic unit tests for ``DataDesignerResource``. + +These tests cover the SDK seams that don't need a server: +- URL construction (``sdk_http.url``) +- The streaming-preview state machine (``_PreviewFrameCollector``) and its + decode helpers (``_decode_preview_frame`` / ``_parse_preview_payload``) +- The client-side seed-source validation gate (``_get_config_for_api_call``) +- The HTTP-error to typed-exception translator (``_get_error``) + +Round-trip tests against the real preview/jobs endpoints live in +``tests/integration/`` (``test_preview.py``, ``test_jobs.py``, +``test_validation.py``, ``test_model_providers.py``). +""" + from collections.abc import AsyncIterator -from datetime import datetime from typing import Any, cast from unittest.mock import patch @@ -11,7 +23,6 @@ import httpx import pandas as pd import pytest -import respx from data_designer.config.analysis.column_statistics import GeneralColumnStatistics from data_designer.config.analysis.dataset_profiler import DatasetProfilerResults from data_designer.config.dataset_metadata import DatasetMetadata @@ -28,10 +39,12 @@ DataDesignerConfigValidationError, DataDesignerPreviewError, ) -from nemo_data_designer_plugin.sdk.job_resources import AsyncDataDesignerJobResource, DataDesignerJobResource -from nemo_data_designer_plugin.sdk.resources import AsyncDataDesignerResource, DataDesignerResource +from nemo_data_designer_plugin.sdk.resources import ( + AsyncDataDesignerResource, + DataDesignerResource, + _decode_preview_frame, +) from nemo_platform import AsyncNeMoPlatform, NeMoPlatform -from nemo_platform.types.inference import ModelProvider as NMPModelProvider from nemo_platform_plugin.functions.frames import Done, Error, Heartbeat from pydantic import BaseModel @@ -69,7 +82,7 @@ def config_builder() -> dd.DataDesignerConfigBuilder: return builder -def make_basic_dataset() -> pd.DataFrame: +def _make_basic_dataset() -> pd.DataFrame: return pd.DataFrame(data={"foo": [1, 2, 3]}).convert_dtypes(dtype_backend="pyarrow") @@ -78,8 +91,8 @@ async def _async_iter(frames: list[BaseModel]) -> AsyncIterator[BaseModel]: yield frame -def make_successful_preview_frames() -> list[BaseModel]: - dataset = make_basic_dataset() +def _make_successful_preview_frames() -> list[BaseModel]: + dataset = _make_basic_dataset() dataset_dict = cast(list[dict[str, Any]], dataset.to_dict(orient="records")) dataset_metadata = DatasetMetadata() analysis = DatasetProfilerResults( @@ -108,6 +121,11 @@ def make_successful_preview_frames() -> list[BaseModel]: ] +# --------------------------------------------------------------------------- +# URL construction +# --------------------------------------------------------------------------- + + @pytest.mark.parametrize("path", ["preview", "/preview", "///preview"]) def test_http_url_normalizes_leading_slashes(platform: NeMoPlatform, path: str) -> None: assert sdk_http.url(platform, None, path) == "http://testserver/apis/data-designer/v2/workspaces/default/preview" @@ -117,27 +135,98 @@ def test_http_url_normalizes_empty_path(platform: NeMoPlatform) -> None: assert sdk_http.url(platform, None, "") == "http://testserver/apis/data-designer/v2/workspaces/default/" -def test_preview_success(resource: DataDesignerResource, config_builder: dd.DataDesignerConfigBuilder) -> None: - with patch.object(resource, "_preview", return_value=make_successful_preview_frames()): +# --------------------------------------------------------------------------- +# Preview frame decoding +# --------------------------------------------------------------------------- + + +def test_decode_preview_frame_returns_typed_frame_for_known_kind() -> None: + frame = _decode_preview_frame('{"kind":"log","level":"info","message":"hello"}') + assert isinstance(frame, LogFrame) + assert frame.message == "hello" + + +def test_decode_preview_frame_returns_none_for_unknown_kind() -> None: + """Unknown frame kinds are forward-compat: dropped silently so older clients can talk to + newer servers without crashing.""" + assert _decode_preview_frame('{"kind":"future","payload":1}') is None + + +def test_decode_preview_frame_returns_none_for_non_object_payload() -> None: + """Defensive: a JSON scalar / list isn't a frame and shouldn't crash the parser.""" + assert _decode_preview_frame("[]") is None + assert _decode_preview_frame('"oops"') is None + + +# --------------------------------------------------------------------------- +# _PreviewFrameCollector behavior (driven via patched _preview) +# --------------------------------------------------------------------------- + + +def test_preview_collector_assembles_dataset_metadata_and_processor_artifacts( + resource: DataDesignerResource, config_builder: dd.DataDesignerConfigBuilder +) -> None: + with patch.object(resource, "_preview", return_value=_make_successful_preview_frames()): preview_results = resource.preview(config_builder) assert isinstance(preview_results.dataset, pd.DataFrame) - pd.testing.assert_frame_equal(preview_results.dataset, make_basic_dataset()) + pd.testing.assert_frame_equal(preview_results.dataset, _make_basic_dataset()) assert preview_results.processor_artifacts == {"processor": [{"foo": "bar"}]} +@pytest.mark.asyncio +async def test_preview_collector_assembles_dataset_async( + async_resource: AsyncDataDesignerResource, config_builder: dd.DataDesignerConfigBuilder +) -> None: + with patch.object(async_resource, "_preview", return_value=_async_iter(_make_successful_preview_frames())): + preview_results = await async_resource.preview(config_builder) + + assert isinstance(preview_results.dataset, pd.DataFrame) + pd.testing.assert_frame_equal(preview_results.dataset, _make_basic_dataset()) + + +def test_preview_collector_raises_when_dataset_frame_is_empty( + resource: DataDesignerResource, config_builder: dd.DataDesignerConfigBuilder +) -> None: + """``_PreviewFrameCollector._accept_dataset`` rejects empty record batches because the + real failure mode is silent column-generation failures, not legitimate empty datasets. + """ + with patch.object(resource, "_preview", return_value=[DatasetFrame(records=[])]): + with pytest.raises(DataDesignerPreviewError): + resource.preview(config_builder) + + +def test_preview_collector_propagates_error_frame_message( + resource: DataDesignerResource, config_builder: dd.DataDesignerConfigBuilder +) -> None: + with patch.object(resource, "_preview", return_value=[Error(message="boom")]): + with pytest.raises(DataDesignerPreviewError, match="boom"): + resource.preview(config_builder) + + +# --------------------------------------------------------------------------- +# Client-side seed-source validation gate (_get_config_for_api_call) +# --------------------------------------------------------------------------- + + @pytest.mark.parametrize("seed_kind", ["df", "local", "directory", "file_contents"]) -def test_preview_rejects_remote_unsupported_seed_sources( +def test_preview_rejects_local_only_seed_sources_before_sending_request( resource: DataDesignerResource, config_builder: dd.DataDesignerConfigBuilder, seed_kind: str, tmp_path, ) -> None: + """The validation gate inside ``_get_config_for_api_call`` rejects seed sources that + only make sense locally (DataFrame, LocalFile, Directory, FileContents), so the SDK + fails fast with a typed error instead of letting the server emit a 422 round-trip + later. Patching ``_preview`` to raise an AssertionError catches any regression where + the request is sent anyway. + """ if seed_kind == "df": seed_source = dd.DataFrameSeedSource(df=pd.DataFrame(data={"foo": [1, 2, 3]})) elif seed_kind == "local": seed_file = tmp_path / "seed.parquet" - make_basic_dataset().to_parquet(seed_file) + _make_basic_dataset().to_parquet(seed_file) seed_source = dd.LocalFileSeedSource(path=str(seed_file)) elif seed_kind == "directory": seed_source = dd.DirectorySeedSource(path=str(tmp_path)) @@ -155,144 +244,52 @@ def test_preview_rejects_remote_unsupported_seed_sources( assert "only supports seed data" in str(exc_info.value) -def test_empty_dataset_frame_raises_preview_error( - resource: DataDesignerResource, config_builder: dd.DataDesignerConfigBuilder -) -> None: - with patch.object(resource, "_preview", return_value=[DatasetFrame(records=[])]): - with pytest.raises(DataDesignerPreviewError): - resource.preview(config_builder) +# --------------------------------------------------------------------------- +# Default model surfaces +# --------------------------------------------------------------------------- -def test_error_frame_raises_preview_error( - resource: DataDesignerResource, config_builder: dd.DataDesignerConfigBuilder -) -> None: - with patch.object(resource, "_preview", return_value=[Error(message="boom")]): - with pytest.raises(DataDesignerPreviewError, match="boom"): - resource.preview(config_builder) - - -@respx.mock -def test_preview_posts_jsonl_request( - resource: DataDesignerResource, config_builder: dd.DataDesignerConfigBuilder -) -> None: - preview_messages = "\n".join(frame.model_dump_json() for frame in make_successful_preview_frames()) + "\n" - route = respx.post("http://testserver/apis/data-designer/v2/workspaces/default/preview").mock( - return_value=httpx.Response(200, text=preview_messages) - ) - - preview_results = resource.preview(config_builder, num_records=3) - - request_json = json.loads(route.calls[0].request.content) - assert request_json["config"]["columns"][0]["column_type"] == "sampler" - assert request_json["num_records"] == 3 - assert preview_results.dataset is not None - - -@respx.mock -def test_preview_ignores_unknown_frame_kind( - resource: DataDesignerResource, config_builder: dd.DataDesignerConfigBuilder -) -> None: - preview_messages = '{"kind":"future","payload":1}\n' + "\n".join( - frame.model_dump_json() for frame in make_successful_preview_frames() - ) - respx.post("http://testserver/apis/data-designer/v2/workspaces/default/preview").mock( - return_value=httpx.Response(200, text=preview_messages) - ) - - preview_results = resource.preview(config_builder, num_records=3) - - assert preview_results.dataset is not None - - -@respx.mock -def test_create_job(resource: DataDesignerResource, config_builder: dd.DataDesignerConfigBuilder) -> None: - route = respx.post("http://testserver/apis/data-designer/v2/workspaces/default/jobs/create").mock( - return_value=httpx.Response(200, json={"name": "data-designer-abc123"}) - ) - - job_resource = resource.create(config_builder) - - assert isinstance(job_resource, DataDesignerJobResource) - request_json = json.loads(route.calls[0].request.content) - assert request_json["spec"]["config"]["columns"][0]["column_type"] == "sampler" - - -@respx.mock -def test_get_job_resource(resource: DataDesignerResource) -> None: - respx.get("http://testserver/apis/data-designer/v2/workspaces/default/jobs/create/data-designer-abc123").mock( - return_value=httpx.Response(200, json={"name": "data-designer-abc123"}) - ) - - job_resource = resource.get_job_resource("data-designer-abc123") - assert isinstance(job_resource, DataDesignerJobResource) - - -def test_get_default_model_configs(resource: DataDesignerResource) -> None: +def test_get_default_model_configs_returns_empty_list(resource: DataDesignerResource) -> None: + """Default model configs aren't supported on the NeMo Platform; the resource returns an + empty list rather than raising, so callers can still build a config without LLMs.""" assert resource.get_default_model_configs() == [] -def test_get_default_model_providers(platform: NeMoPlatform, resource: DataDesignerResource) -> None: - mock_providers = [ - NMPModelProvider( - name="provider1", - workspace="ws1", - host_url="http://host1", - created_at=datetime.now(), - updated_at=datetime.now(), - ), - NMPModelProvider( - name="provider2", - workspace="ws2", - host_url="http://host2", - created_at=datetime.now(), - updated_at=datetime.now(), - ), - ] - with patch.object(platform.inference.providers, "list", return_value=mock_providers): - providers = resource.get_default_model_providers() - assert len(providers) == 2 - assert all(isinstance(provider, dd.ModelProvider) for provider in providers) +# --------------------------------------------------------------------------- +# HTTP-error translation (_get_error) +# --------------------------------------------------------------------------- @pytest.mark.asyncio -async def test_preview_success_async( +async def test_http_error_translates_422_to_config_validation_error( async_resource: AsyncDataDesignerResource, config_builder: dd.DataDesignerConfigBuilder ) -> None: - with patch.object(async_resource, "_preview", return_value=_async_iter(make_successful_preview_frames())): - preview_results = await async_resource.preview(config_builder) - - assert isinstance(preview_results.dataset, pd.DataFrame) - pd.testing.assert_frame_equal(preview_results.dataset, make_basic_dataset()) - + request = httpx.Request("POST", "http://testserver") + response = httpx.Response(422, request=request, json={"detail": "bad config"}) -@pytest.mark.asyncio -@respx.mock -async def test_create_job_async( - async_resource: AsyncDataDesignerResource, config_builder: dd.DataDesignerConfigBuilder -) -> None: - respx.post("http://testserver/apis/data-designer/v2/workspaces/default/jobs/create").mock( - return_value=httpx.Response(200, json={"name": "data-designer-abc123"}) - ) + with patch.object( + async_resource, "_preview", side_effect=httpx.HTTPStatusError("bad", request=request, response=response) + ): + with pytest.raises(DataDesignerConfigValidationError) as exc_info: + await async_resource.preview(config_builder) - job_resource = await async_resource.create(config_builder) - assert isinstance(job_resource, AsyncDataDesignerJobResource) + assert exc_info.value.status_code == 422 + assert "bad config" in str(exc_info.value) @pytest.mark.asyncio -async def test_http_error_handling_async( +async def test_http_error_translates_5xx_to_generic_client_error( async_resource: AsyncDataDesignerResource, config_builder: dd.DataDesignerConfigBuilder ) -> None: request = httpx.Request("POST", "http://testserver") - response = httpx.Response(422, request=request, json={"detail": "bad config"}) - with patch.object( - async_resource, "_preview", side_effect=httpx.HTTPStatusError("bad", request=request, response=response) - ): - with pytest.raises(DataDesignerConfigValidationError): - await async_resource.preview(config_builder) - response = httpx.Response(500, request=request, text="boom") + with patch.object( async_resource, "_preview", side_effect=httpx.HTTPStatusError("bad", request=request, response=response) ): - with pytest.raises(DataDesignerClientError): + with pytest.raises(DataDesignerClientError) as exc_info: await async_resource.preview(config_builder) + + assert exc_info.value.status_code == 500 + # 5xx is *not* a config validation error — make sure we didn't accidentally widen the 422 branch. + assert not isinstance(exc_info.value, DataDesignerConfigValidationError)