chore: refactor data-designer plugin tests

plugins/nemo-data-designer/src/nemo_data_designer_plugin/sdk/errors.py

@@ -1,11 +1,23 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+from __future__ import annotations
+
+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 +30,27 @@ 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
+    try:
+        body = response.json()
+    except Exception:
+        body = None
+    if isinstance(body, dict) and isinstance(body.get("detail"), str):
+        detail = body["detail"]
+
+    return response.status_code, detail

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:

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}")
 
 

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(