From 7f053fc7779317e0bddde824e59dc20b444ef1b0 Mon Sep 17 00:00:00 2001 From: Mike Knepper Date: Wed, 20 May 2026 11:55:02 -0500 Subject: [PATCH 01/17] Move personas CLI tests to integration dir Signed-off-by: Mike Knepper --- .../tests/integration/test_personas_cli.py | 226 ++++++++++++++++++ .../nemo-data-designer/tests/unit/test_cli.py | 190 +-------------- 2 files changed, 228 insertions(+), 188 deletions(-) create mode 100644 plugins/nemo-data-designer/tests/integration/test_personas_cli.py diff --git a/plugins/nemo-data-designer/tests/integration/test_personas_cli.py b/plugins/nemo-data-designer/tests/integration/test_personas_cli.py new file mode 100644 index 0000000000..c46d998315 --- /dev/null +++ b/plugins/nemo-data-designer/tests/integration/test_personas_cli.py @@ -0,0 +1,226 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +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_platform import NeMoPlatform +from nemo_platform.types.files import NGCStorageConfig +from typer.testing import CliRunner + + +@pytest.fixture +def mock_ngc_client() -> Generator[dict[str, Mock]]: + with ( + patch("nmp.core.files.app.backends.ngc.Client") as mock_client_cls, + patch("nmp.core.files.app.backends.ngc.ResourceAPI") as mock_resource_api_cls, + ): + mock_client = Mock() + mock_resource_api = Mock() + + mock_client_cls.return_value = mock_client + mock_resource_api_cls.return_value = mock_resource_api + + yield { + "client": mock_client, + "resource_api": mock_resource_api, + } + + +@pytest.fixture +def sdk(monkeypatch: pytest.MonkeyPatch, mock_ngc_client: dict[str, Mock]) -> Generator[NeMoPlatform]: + with u.make_mock_client_context() as client_context: + monkeypatch.setenv("NGC_API_KEY", "nvapi-abc123") + yield client_context.sdk + monkeypatch.delenv("NGC_API_KEY") + + +@pytest.fixture +def runner() -> CliRunner: + return CliRunner() + + +@pytest.fixture +def app() -> typer.Typer: + return u.make_data_designer_cli_app() + + +@pytest.fixture +def cli_sdk(monkeypatch: pytest.MonkeyPatch, sdk: NeMoPlatform) -> NeMoPlatform: + monkeypatch.setattr(personas_module, "NeMoPlatform", lambda: sdk) + return sdk + + +@pytest.mark.integration +def test_make_fileset_creates_requested_locale_with_existing_secret( + runner: CliRunner, app: typer.Typer, cli_sdk: NeMoPlatform +) -> None: + result = runner.invoke( + app, + [ + "personas", + "make-fileset", + "--locale", + "en_US", + "--api-key-secret", + "system/ngc-api-key", + ], + ) + + assert result.exit_code == 0, result.output + filesets = cli_sdk.files.filesets.list(workspace=WORKSPACE) + assert [fileset.name for fileset in filesets.data] == [get_resource_name_for_locale("en_US")] + + fileset = cli_sdk.files.filesets.retrieve(name=get_resource_name_for_locale("en_US"), workspace=WORKSPACE) + assert isinstance(fileset.storage, NGCStorageConfig) + assert fileset.storage.api_key_secret == "system/ngc-api-key" + + +@pytest.mark.integration +def test_make_fileset_creates_secret_from_env_then_fileset( + monkeypatch: pytest.MonkeyPatch, runner: CliRunner, app: typer.Typer, cli_sdk: NeMoPlatform +) -> None: + monkeypatch.setenv("MY_NGC_API_KEY", "nvapi-from-env") + + result = runner.invoke( + app, + [ + "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 + secret = cli_sdk.secrets.access("my-ngc-key", workspace="system") + assert secret.value == "nvapi-from-env" + + fileset = cli_sdk.files.filesets.retrieve(name=get_resource_name_for_locale("en_US"), workspace=WORKSPACE) + assert isinstance(fileset.storage, NGCStorageConfig) + assert fileset.storage.api_key_secret == "system/my-ngc-key" + + +@pytest.mark.integration +def test_make_fileset_missing_env_var_is_clear(runner: CliRunner, app: typer.Typer) -> None: + result = runner.invoke( + app, + [ + "personas", + "make-fileset", + "--locale", + "en_US", + "--api-key-secret", + "system/my-ngc-key", + "--api-key-env-var", + "MISSING_NGC_API_KEY", + ], + ) + + assert result.exit_code != 0 + assert "MISSING_NGC_API_KEY" in result.output + assert "not set or is empty" in result.output + + +@pytest.mark.integration +def test_make_fileset_unknown_locale_is_clear(runner: CliRunner, app: typer.Typer) -> None: + result = runner.invoke( + app, + [ + "personas", + "make-fileset", + "--locale", + "de_DE", + "--api-key-secret", + "system/ngc-api-key", + ], + ) + + assert result.exit_code != 0 + assert "Invalid value for '--locale'" in result.output + assert "de_DE" in result.output + + +@pytest.mark.integration +def test_make_fileset_bare_secret_name_is_clear(runner: CliRunner, app: typer.Typer) -> None: + result = runner.invoke( + app, + [ + "personas", + "make-fileset", + "--locale", + "en_US", + "--api-key-secret", + "ngc-api-key", + ], + ) + + assert result.exit_code != 0 + assert "WORKSPACE/NAME" in result.output + + +@pytest.mark.integration +def test_make_fileset_create_secret_conflict_does_not_create_fileset( + monkeypatch: pytest.MonkeyPatch, runner: CliRunner, app: typer.Typer, 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, + [ + "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 == 1 + assert "already exists" in result.output + filesets = cli_sdk.files.filesets.list(workspace=WORKSPACE) + assert filesets.data == [] + + +@pytest.mark.integration +def test_make_fileset_create_secret_internal_error_surfaces_clearly( + monkeypatch: pytest.MonkeyPatch, runner: CliRunner, app: typer.Typer, cli_sdk: NeMoPlatform +) -> None: + monkeypatch.setenv("MY_NGC_API_KEY", "nvapi-from-env") + + def _boom(*args: object, **kwargs: object) -> None: + raise RuntimeError("secrets backend exploded") + + with patch.object(cli_sdk.secrets, "create", side_effect=_boom): + result = runner.invoke( + app, + [ + "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 == 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/unit/test_cli.py b/plugins/nemo-data-designer/tests/unit/test_cli.py index 4bb6eb73c1..5f9bd6f68c 100644 --- a/plugins/nemo-data-designer/tests/unit/test_cli.py +++ b/plugins/nemo-data-designer/tests/unit/test_cli.py @@ -1,55 +1,12 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -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 -@pytest.fixture -def mock_ngc_client() -> Generator[dict[str, Mock]]: - with ( - patch("nmp.core.files.app.backends.ngc.Client") as mock_client_cls, - patch("nmp.core.files.app.backends.ngc.ResourceAPI") as mock_resource_api_cls, - ): - mock_client = Mock() - mock_resource_api = Mock() - - mock_client_cls.return_value = mock_client - mock_resource_api_cls.return_value = mock_resource_api - - yield { - "client": mock_client, - "resource_api": mock_resource_api, - } - - -@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: - monkeypatch.setenv("NGC_API_KEY", "nvapi-abc123") - yield sdk - monkeypatch.delenv("NGC_API_KEY") - - @pytest.fixture def runner() -> CliRunner: return CliRunner() @@ -57,150 +14,7 @@ def runner() -> 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, - [ - "personas", - "make-fileset", - "--locale", - "en_US", - "--api-key-secret", - "system/ngc-api-key", - ], - ) - - assert result.exit_code == 0, result.output - filesets = cli_sdk.files.filesets.list(workspace=WORKSPACE) - assert [fileset.name for fileset in filesets.data] == [get_resource_name_for_locale("en_US")] - - fileset = cli_sdk.files.filesets.retrieve(name=get_resource_name_for_locale("en_US"), workspace=WORKSPACE) - assert isinstance(fileset.storage, NGCStorageConfig) - assert fileset.storage.api_key_secret == "system/ngc-api-key" - - -def test_make_fileset_creates_secret_from_env_then_fileset( - monkeypatch: pytest.MonkeyPatch, runner: CliRunner, app: typer.Typer, cli_sdk: NeMoPlatform -) -> None: - monkeypatch.setenv("MY_NGC_API_KEY", "nvapi-from-env") - - result = runner.invoke( - app, - [ - "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 - secret = cli_sdk.secrets.access("my-ngc-key", workspace="system") - assert secret.value == "nvapi-from-env" - - fileset = cli_sdk.files.filesets.retrieve(name=get_resource_name_for_locale("en_US"), workspace=WORKSPACE) - assert isinstance(fileset.storage, NGCStorageConfig) - 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, - [ - "personas", - "make-fileset", - "--locale", - "en_US", - "--api-key-secret", - "system/my-ngc-key", - "--api-key-env-var", - "MISSING_NGC_API_KEY", - ], - ) - - assert result.exit_code != 0 - assert "MISSING_NGC_API_KEY" in result.output - 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, - [ - "personas", - "make-fileset", - "--locale", - "de_DE", - "--api-key-secret", - "system/ngc-api-key", - ], - ) - - assert result.exit_code != 0 - assert "Invalid value for '--locale'" in result.output - 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, - [ - "personas", - "make-fileset", - "--locale", - "en_US", - "--api-key-secret", - "ngc-api-key", - ], - ) - - assert result.exit_code != 0 - assert "WORKSPACE/NAME" in result.output - - -def test_make_fileset_create_secret_conflict_does_not_create_fileset( - monkeypatch: pytest.MonkeyPatch, runner: CliRunner, app: typer.Typer, 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, - [ - "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 == 1 - assert "already exists" in result.output - filesets = cli_sdk.files.filesets.list(workspace=WORKSPACE) - assert filesets.data == [] + return u.make_data_designer_cli_app() def test_nemotron_personas_download_is_wired(runner: CliRunner, app: typer.Typer) -> None: From 6ccced74c6401a5af1d48605b3cb9e5c9c50e3be Mon Sep 17 00:00:00 2001 From: Mike Knepper Date: Wed, 20 May 2026 13:14:17 -0500 Subject: [PATCH 02/17] Clean up model provider tests Signed-off-by: Mike Knepper --- .../tests/integration/test_model_providers.py | 36 +++++ .../tests/integration/test_validation.py | 54 ++++--- .../tests/unit/test_model_provider.py | 145 ++---------------- 3 files changed, 85 insertions(+), 150 deletions(-) create mode 100644 plugins/nemo-data-designer/tests/integration/test_model_providers.py diff --git a/plugins/nemo-data-designer/tests/integration/test_model_providers.py b/plugins/nemo-data-designer/tests/integration/test_model_providers.py new file mode 100644 index 0000000000..15da7f64ed --- /dev/null +++ b/plugins/nemo-data-designer/tests/integration/test_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 + + +@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" + + +@pytest.mark.integration +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/integration/test_validation.py b/plugins/nemo-data-designer/tests/integration/test_validation.py index 56298e2f40..7c5f42c3a4 100644 --- a/plugins/nemo-data-designer/tests/integration/test_validation.py +++ b/plugins/nemo-data-designer/tests/integration/test_validation.py @@ -25,12 +25,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 +34,51 @@ 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 + + +@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 = _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]) + + +@pytest.mark.integration +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]) + + @pytest.mark.integration 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, 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 From e99e15dfd407eb0cb61660f5258bf985c01233ed Mon Sep 17 00:00:00 2001 From: Mike Knepper Date: Wed, 20 May 2026 15:50:36 -0500 Subject: [PATCH 03/17] Job integration test cleanup Signed-off-by: Mike Knepper --- .../testing/utils.py | 5 +- .../tests/integration/test_jobs.py | 367 ++++++++++++++++++ .../tests/unit/test_job_results.py | 57 --- .../tests/unit/test_sdk_job_resources.py | 159 -------- 4 files changed, 371 insertions(+), 217 deletions(-) create mode 100644 plugins/nemo-data-designer/tests/integration/test_jobs.py delete mode 100644 plugins/nemo-data-designer/tests/unit/test_job_results.py delete mode 100644 plugins/nemo-data-designer/tests/unit/test_sdk_job_resources.py 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..7abfc7266b 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 @@ -405,7 +405,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_jobs.py b/plugins/nemo-data-designer/tests/integration/test_jobs.py new file mode 100644 index 0000000000..6fef3634df --- /dev/null +++ b/plugins/nemo-data-designer/tests/integration/test_jobs.py @@ -0,0 +1,367 @@ +# 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 + +_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.integration +@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.integration +@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.integration +@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.integration +@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.integration +@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.integration +@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.integration +@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.integration +@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. Pagination is therefore covered indirectly through +# ``wait_until_done``'s ``_poll_safe(self.get_logs, ...)`` calls. +# --------------------------------------------------------------------------- + + +# --------------------------------------------------------------------------- +# download_artifacts (sync + async) and DataDesignerJobResults +# --------------------------------------------------------------------------- + + +@pytest.mark.integration +@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.integration +@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.integration +@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.integration +@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.integration +@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.integration +@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.integration +@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.integration +@pytest.mark.asyncio +async def test_load_analysis_raises_when_active_but_result_missing() -> None: + """An ``active`` job whose analysis result has not yet been written returns a 404 from the + Jobs service; the resource raises ``DataDesignerJobError``. + + Note: ``_check_if_result_available`` has a friendly-message branch for 404s + (``f"{result_name!r} result is not available."``), but that branch checks for + the literal substring ``"404"`` in the error message, while the real Jobs service + returns ``"Job result not found"`` with no status code in the body. The user + therefore gets the generic "Error loading dataset" message instead. We assert + the current behavior; the friendly-message detection logic is worth a follow-up. + """ + 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="Job result not found"): + job_resource.load_analysis() 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_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) From 36afa0248c41564429291dd4c677d362a2ef7d75 Mon Sep 17 00:00:00 2001 From: Mike Knepper Date: Wed, 20 May 2026 15:59:16 -0500 Subject: [PATCH 04/17] Better client-side error parsing Signed-off-by: Mike Knepper --- .../nemo_data_designer_plugin/sdk/errors.py | 41 ++++++++- .../sdk/job_resources.py | 10 +-- .../sdk/resources.py | 19 ++-- .../tests/integration/test_jobs.py | 17 ++-- .../tests/unit/test_errors.py | 87 +++++++++++++++++++ 5 files changed, 143 insertions(+), 31 deletions(-) create mode 100644 plugins/nemo-data-designer/tests/unit/test_errors.py 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..6fde011ea0 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,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,30 @@ 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. We do this because the basic httpx client we + use in this SDK doesn't expose the structured error types that the + Stainless-generated client did, so the response body is the only source of + truth for the human-facing message. + """ + 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 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/tests/integration/test_jobs.py b/plugins/nemo-data-designer/tests/integration/test_jobs.py index 6fef3634df..c2d3f990b7 100644 --- a/plugins/nemo-data-designer/tests/integration/test_jobs.py +++ b/plugins/nemo-data-designer/tests/integration/test_jobs.py @@ -348,20 +348,15 @@ async def test_load_analysis_when_terminally_incomplete_warns_and_returns_partia @pytest.mark.integration @pytest.mark.asyncio -async def test_load_analysis_raises_when_active_but_result_missing() -> None: - """An ``active`` job whose analysis result has not yet been written returns a 404 from the - Jobs service; the resource raises ``DataDesignerJobError``. - - Note: ``_check_if_result_available`` has a friendly-message branch for 404s - (``f"{result_name!r} result is not available."``), but that branch checks for - the literal substring ``"404"`` in the error message, while the real Jobs service - returns ``"Job result not found"`` with no status code in the body. The user - therefore gets the generic "Error loading dataset" message instead. We assert - the current behavior; the friendly-message detection logic is worth a follow-up. +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="Job result not found"): + with pytest.raises(DataDesignerJobError, match="'analysis' result is not available"): job_resource.load_analysis() 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 From fbb2c5369701c161cb6f13cfefa5a4cb8ddebeb3 Mon Sep 17 00:00:00 2001 From: Mike Knepper Date: Wed, 20 May 2026 16:18:16 -0500 Subject: [PATCH 05/17] Trim sdk resources unit tests, most covered by integration tests Signed-off-by: Mike Knepper --- .../tests/unit/test_sdk_resources.py | 261 +++++++++--------- 1 file changed, 129 insertions(+), 132 deletions(-) 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) From 3e96e51c0e6b40b675bfac4152b24a0daa22f9ac Mon Sep 17 00:00:00 2001 From: Mike Knepper Date: Wed, 20 May 2026 16:41:32 -0500 Subject: [PATCH 06/17] Test preview error through the sdk Signed-off-by: Mike Knepper --- .../integration/test_preview_streaming.py | 43 +++++++++++++++++++ .../tests/unit/test_preview_function.py | 23 ---------- 2 files changed, 43 insertions(+), 23 deletions(-) create mode 100644 plugins/nemo-data-designer/tests/integration/test_preview_streaming.py diff --git a/plugins/nemo-data-designer/tests/integration/test_preview_streaming.py b/plugins/nemo-data-designer/tests/integration/test_preview_streaming.py new file mode 100644 index 0000000000..644aebe66d --- /dev/null +++ b/plugins/nemo-data-designer/tests/integration/test_preview_streaming.py @@ -0,0 +1,43 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""End-to-end coverage of preview-stream error propagation. + +This complements ``test_preview.py`` (happy-path streaming) by exercising +the worker-raised-an-exception path all the way from the in-process FastAPI +route through the NDJSON stream and SDK frame decoder back into a typed +``DataDesignerPreviewError``. +""" + +import data_designer.config as dd +import nemo_data_designer_plugin.testing.utils as u +import pytest +from nemo_data_designer_plugin.functions import _preview_worker as worker_module +from nemo_data_designer_plugin.sdk.errors import DataDesignerPreviewError + + +@pytest.mark.integration +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. + """ + + 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) 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) From de50fcf64363701b1bc2c10bed611a2229d68367 Mon Sep 17 00:00:00 2001 From: Mike Knepper Date: Thu, 21 May 2026 09:09:59 -0500 Subject: [PATCH 07/17] docstring cleanup Signed-off-by: Mike Knepper --- .../src/nemo_data_designer_plugin/sdk/errors.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) 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 6fde011ea0..60e02cfd26 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 @@ -37,10 +37,7 @@ def extract_http_error_info(exc: httpx.HTTPStatusError) -> tuple[int, str]: 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. We do this because the basic httpx client we - use in this SDK doesn't expose the structured error types that the - Stainless-generated client did, so the response body is the only source of - truth for the human-facing message. + text if that isn't available. """ response = exc.response try: From a1332f9320e989f4f9dc17f7c652829aed149dde Mon Sep 17 00:00:00 2001 From: Mike Knepper Date: Thu, 21 May 2026 09:41:02 -0500 Subject: [PATCH 08/17] consolidate preview tests Signed-off-by: Mike Knepper --- .../tests/integration/test_preview.py | 32 +++++++++++++- .../integration/test_preview_streaming.py | 43 ------------------- 2 files changed, 30 insertions(+), 45 deletions(-) delete mode 100644 plugins/nemo-data-designer/tests/integration/test_preview_streaming.py diff --git a/plugins/nemo-data-designer/tests/integration/test_preview.py b/plugins/nemo-data-designer/tests/integration/test_preview.py index 48e0ad8842..c535e740b4 100644 --- a/plugins/nemo-data-designer/tests/integration/test_preview.py +++ b/plugins/nemo-data-designer/tests/integration/test_preview.py @@ -11,7 +11,7 @@ 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 @pytest.mark.integration @@ -68,7 +68,7 @@ def test_happy_path_preview() -> None: @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) @@ -175,6 +175,34 @@ def test_preview_with_schema_transform_processor() -> None: assert "messages" in processor_records[0] +@pytest.mark.integration +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_preview_streaming.py b/plugins/nemo-data-designer/tests/integration/test_preview_streaming.py deleted file mode 100644 index 644aebe66d..0000000000 --- a/plugins/nemo-data-designer/tests/integration/test_preview_streaming.py +++ /dev/null @@ -1,43 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""End-to-end coverage of preview-stream error propagation. - -This complements ``test_preview.py`` (happy-path streaming) by exercising -the worker-raised-an-exception path all the way from the in-process FastAPI -route through the NDJSON stream and SDK frame decoder back into a typed -``DataDesignerPreviewError``. -""" - -import data_designer.config as dd -import nemo_data_designer_plugin.testing.utils as u -import pytest -from nemo_data_designer_plugin.functions import _preview_worker as worker_module -from nemo_data_designer_plugin.sdk.errors import DataDesignerPreviewError - - -@pytest.mark.integration -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. - """ - - 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) From f5955f5dce07a195bb5cf825e17663ef9688f8bc Mon Sep 17 00:00:00 2001 From: Mike Knepper Date: Thu, 21 May 2026 10:29:05 -0500 Subject: [PATCH 09/17] Refactor to shared invoke_cli helper Signed-off-by: Mike Knepper --- .../testing/utils.py | 25 +++++++- .../tests/integration/test_cli_local.py | 43 ++++--------- .../tests/integration/test_personas_cli.py | 63 +++++++------------ .../nemo-data-designer/tests/unit/test_cli.py | 20 ++---- 4 files changed, 60 insertions(+), 91 deletions(-) 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 7abfc7266b..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") diff --git a/plugins/nemo-data-designer/tests/integration/test_cli_local.py b/plugins/nemo-data-designer/tests/integration/test_cli_local.py index 326fe1b85b..e255b0a2d9 100644 --- a/plugins/nemo-data-designer/tests/integration/test_cli_local.py +++ b/plugins/nemo-data-designer/tests/integration/test_cli_local.py @@ -6,30 +6,17 @@ 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 - - -@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 +27,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 @@ -54,7 +41,7 @@ def test_preview_run_saves_expected_artifacts(runner: CliRunner, app: typer.Type @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 @@ -95,9 +81,7 @@ def load_config_builder() -> dd.DataDesignerConfigBuilder: @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 +99,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 @@ -131,14 +114,14 @@ def load_config_builder() -> dd.DataDesignerConfigBuilder: @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_personas_cli.py b/plugins/nemo-data-designer/tests/integration/test_personas_cli.py index c46d998315..6b7f5f749e 100644 --- a/plugins/nemo-data-designer/tests/integration/test_personas_cli.py +++ b/plugins/nemo-data-designer/tests/integration/test_personas_cli.py @@ -6,12 +6,10 @@ 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_platform import NeMoPlatform from nemo_platform.types.files import NGCStorageConfig -from typer.testing import CliRunner @pytest.fixture @@ -40,16 +38,6 @@ def sdk(monkeypatch: pytest.MonkeyPatch, mock_ngc_client: dict[str, Mock]) -> Ge monkeypatch.delenv("NGC_API_KEY") -@pytest.fixture -def runner() -> CliRunner: - return CliRunner() - - -@pytest.fixture -def app() -> typer.Typer: - return u.make_data_designer_cli_app() - - @pytest.fixture def cli_sdk(monkeypatch: pytest.MonkeyPatch, sdk: NeMoPlatform) -> NeMoPlatform: monkeypatch.setattr(personas_module, "NeMoPlatform", lambda: sdk) @@ -57,11 +45,8 @@ def cli_sdk(monkeypatch: pytest.MonkeyPatch, sdk: NeMoPlatform) -> NeMoPlatform: @pytest.mark.integration -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_make_fileset_creates_requested_locale_with_existing_secret(cli_sdk: NeMoPlatform) -> None: + result = u.invoke_cli( [ "personas", "make-fileset", @@ -69,7 +54,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 @@ -83,12 +68,11 @@ def test_make_fileset_creates_requested_locale_with_existing_secret( @pytest.mark.integration 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", @@ -98,7 +82,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 @@ -111,9 +95,8 @@ def test_make_fileset_creates_secret_from_env_then_fileset( @pytest.mark.integration -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", @@ -123,7 +106,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 @@ -132,9 +115,8 @@ def test_make_fileset_missing_env_var_is_clear(runner: CliRunner, app: typer.Typ @pytest.mark.integration -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", @@ -142,7 +124,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 @@ -151,9 +133,8 @@ def test_make_fileset_unknown_locale_is_clear(runner: CliRunner, app: typer.Type @pytest.mark.integration -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", @@ -161,7 +142,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 @@ -170,13 +151,12 @@ def test_make_fileset_bare_secret_name_is_clear(runner: CliRunner, app: typer.Ty @pytest.mark.integration 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", @@ -186,7 +166,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 @@ -197,7 +177,7 @@ def test_make_fileset_create_secret_conflict_does_not_create_fileset( @pytest.mark.integration def test_make_fileset_create_secret_internal_error_surfaces_clearly( - 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") @@ -205,8 +185,7 @@ def _boom(*args: object, **kwargs: object) -> None: raise RuntimeError("secrets backend exploded") with patch.object(cli_sdk.secrets, "create", side_effect=_boom): - result = runner.invoke( - app, + result = u.invoke_cli( [ "personas", "make-fileset", @@ -216,7 +195,7 @@ def _boom(*args: object, **kwargs: object) -> None: "system/my-ngc-key", "--api-key-env-var", "MY_NGC_API_KEY", - ], + ] ) assert result.exit_code == 1 diff --git a/plugins/nemo-data-designer/tests/unit/test_cli.py b/plugins/nemo-data-designer/tests/unit/test_cli.py index 5f9bd6f68c..e961ff2ba3 100644 --- a/plugins/nemo-data-designer/tests/unit/test_cli.py +++ b/plugins/nemo-data-designer/tests/unit/test_cli.py @@ -3,22 +3,10 @@ import nemo_data_designer_plugin.testing.utils as u import pytest -import typer -from typer.testing import CliRunner -@pytest.fixture -def runner() -> CliRunner: - return CliRunner() - - -@pytest.fixture -def app() -> typer.Typer: - return u.make_data_designer_cli_app() - - -def test_nemotron_personas_download_is_wired(runner: CliRunner, app: typer.Typer) -> None: - result = runner.invoke(app, ["personas", "download", "--help"]) +def test_nemotron_personas_download_is_wired() -> None: + result = u.invoke_cli(["personas", "download", "--help"]) assert result.exit_code == 0, result.output assert "Download Nemotron-Personas" in result.output @@ -27,8 +15,8 @@ def test_nemotron_personas_download_is_wired(runner: CliRunner, app: typer.Typer @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 test_preview_exposes_save_results_flags(verb: str) -> None: + result = u.invoke_cli(["preview", verb, "--help"]) assert result.exit_code == 0, result.output assert "--save-results" in result.output From b804e6d30427f9eb2cacb5f7c71891e801509dc3 Mon Sep 17 00:00:00 2001 From: Mike Knepper Date: Thu, 21 May 2026 10:45:28 -0500 Subject: [PATCH 10/17] Rename preview-related integration tests Signed-off-by: Mike Knepper --- .../{test_cli_local.py => test_preview_local_cli.py} | 2 +- .../integration/{test_preview.py => test_preview_remote_sdk.py} | 0 2 files changed, 1 insertion(+), 1 deletion(-) rename plugins/nemo-data-designer/tests/integration/{test_cli_local.py => test_preview_local_cli.py} (99%) rename plugins/nemo-data-designer/tests/integration/{test_preview.py => test_preview_remote_sdk.py} (100%) 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 99% 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 e255b0a2d9..8dbb773fd2 100644 --- a/plugins/nemo-data-designer/tests/integration/test_cli_local.py +++ b/plugins/nemo-data-designer/tests/integration/test_preview_local_cli.py @@ -27,7 +27,7 @@ def test_preview_run_saves_expected_artifacts(tmp_path: Path) -> None: "--artifact-path", str(artifact_path), ], - client_context + client_context, ) 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 100% rename from plugins/nemo-data-designer/tests/integration/test_preview.py rename to plugins/nemo-data-designer/tests/integration/test_preview_remote_sdk.py From 89b64f49ab2f7fc282d3e8c10d89fd3acd0f45a4 Mon Sep 17 00:00:00 2001 From: Mike Knepper Date: Thu, 21 May 2026 10:46:44 -0500 Subject: [PATCH 11/17] Move personas download wiring test to other module Signed-off-by: Mike Knepper --- .../tests/integration/test_personas_cli.py | 9 +++++++++ plugins/nemo-data-designer/tests/unit/test_cli.py | 9 --------- 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/plugins/nemo-data-designer/tests/integration/test_personas_cli.py b/plugins/nemo-data-designer/tests/integration/test_personas_cli.py index 6b7f5f749e..535a012a3c 100644 --- a/plugins/nemo-data-designer/tests/integration/test_personas_cli.py +++ b/plugins/nemo-data-designer/tests/integration/test_personas_cli.py @@ -44,6 +44,15 @@ def cli_sdk(monkeypatch: pytest.MonkeyPatch, sdk: NeMoPlatform) -> NeMoPlatform: return sdk +@pytest.mark.integration +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 + + @pytest.mark.integration def test_make_fileset_creates_requested_locale_with_existing_secret(cli_sdk: NeMoPlatform) -> None: result = u.invoke_cli( diff --git a/plugins/nemo-data-designer/tests/unit/test_cli.py b/plugins/nemo-data-designer/tests/unit/test_cli.py index e961ff2ba3..b483572157 100644 --- a/plugins/nemo-data-designer/tests/unit/test_cli.py +++ b/plugins/nemo-data-designer/tests/unit/test_cli.py @@ -5,15 +5,6 @@ import pytest -def test_nemotron_personas_download_is_wired() -> None: - result = u.invoke_cli(["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 - - @pytest.mark.parametrize("verb", ["run", "submit"]) def test_preview_exposes_save_results_flags(verb: str) -> None: result = u.invoke_cli(["preview", verb, "--help"]) From 22b597d0eb5962aeec2b5a43ed7b50c35c32d7d9 Mon Sep 17 00:00:00 2001 From: Mike Knepper Date: Thu, 21 May 2026 10:49:54 -0500 Subject: [PATCH 12/17] Drop worthless unit test Signed-off-by: Mike Knepper --- plugins/nemo-data-designer/tests/unit/test_cli.py | 15 --------------- 1 file changed, 15 deletions(-) delete mode 100644 plugins/nemo-data-designer/tests/unit/test_cli.py diff --git a/plugins/nemo-data-designer/tests/unit/test_cli.py b/plugins/nemo-data-designer/tests/unit/test_cli.py deleted file mode 100644 index b483572157..0000000000 --- a/plugins/nemo-data-designer/tests/unit/test_cli.py +++ /dev/null @@ -1,15 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -import nemo_data_designer_plugin.testing.utils as u -import pytest - - -@pytest.mark.parametrize("verb", ["run", "submit"]) -def test_preview_exposes_save_results_flags(verb: str) -> None: - result = u.invoke_cli(["preview", verb, "--help"]) - - 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 From 86353630a3dedc6c8d40bcc6204d2f42fd1ca9b1 Mon Sep 17 00:00:00 2001 From: Mike Knepper Date: Thu, 21 May 2026 10:54:25 -0500 Subject: [PATCH 13/17] Rename another test file Signed-off-by: Mike Knepper --- .../{test_model_providers.py => test_sdk_get_model_providers.py} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename plugins/nemo-data-designer/tests/integration/{test_model_providers.py => test_sdk_get_model_providers.py} (100%) diff --git a/plugins/nemo-data-designer/tests/integration/test_model_providers.py b/plugins/nemo-data-designer/tests/integration/test_sdk_get_model_providers.py similarity index 100% rename from plugins/nemo-data-designer/tests/integration/test_model_providers.py rename to plugins/nemo-data-designer/tests/integration/test_sdk_get_model_providers.py From e475388ffe2dd41884795b25d1a4d826449f28f7 Mon Sep 17 00:00:00 2001 From: Mike Knepper Date: Thu, 21 May 2026 10:55:41 -0500 Subject: [PATCH 14/17] Rename file Signed-off-by: Mike Knepper --- .../{test_validation.py => test_remote_validation_errors.py} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename plugins/nemo-data-designer/tests/integration/{test_validation.py => test_remote_validation_errors.py} (100%) 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 100% rename from plugins/nemo-data-designer/tests/integration/test_validation.py rename to plugins/nemo-data-designer/tests/integration/test_remote_validation_errors.py From fa8a990216ef118e7913ac9c5552a0d335e947b2 Mon Sep 17 00:00:00 2001 From: Mike Knepper Date: Thu, 21 May 2026 14:33:54 -0500 Subject: [PATCH 15/17] Rename and tweak job-related tests Signed-off-by: Mike Knepper --- .../{test_jobs.py => test_job_sdk.py} | 3 +- .../{test_task.py => test_job_task.py} | 75 +++++++++++++------ 2 files changed, 54 insertions(+), 24 deletions(-) rename plugins/nemo-data-designer/tests/integration/{test_jobs.py => test_job_sdk.py} (99%) rename plugins/nemo-data-designer/tests/integration/{test_task.py => test_job_task.py} (66%) diff --git a/plugins/nemo-data-designer/tests/integration/test_jobs.py b/plugins/nemo-data-designer/tests/integration/test_job_sdk.py similarity index 99% rename from plugins/nemo-data-designer/tests/integration/test_jobs.py rename to plugins/nemo-data-designer/tests/integration/test_job_sdk.py index c2d3f990b7..b69420d8c6 100644 --- a/plugins/nemo-data-designer/tests/integration/test_jobs.py +++ b/plugins/nemo-data-designer/tests/integration/test_job_sdk.py @@ -230,8 +230,7 @@ async def test_wait_until_done_logs_terminal_failure_for_cancelled_status( # 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. Pagination is therefore covered indirectly through -# ``wait_until_done``'s ``_poll_safe(self.get_logs, ...)`` calls. +# the task emitted. # --------------------------------------------------------------------------- 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..7aae501a1c 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,29 @@ 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 + + +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 @@ -40,7 +65,7 @@ def _failing_result_manager() -> Generator[None]: @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,11 +89,13 @@ 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 @@ -77,7 +104,7 @@ async def test_task() -> None: @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,13 @@ 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.integration @pytest.mark.asyncio async def test_exiting_with_error() -> None: builder = dd.DataDesignerConfigBuilder(model_configs=[u.make_model_config()]) @@ -130,21 +160,22 @@ async def test_exiting_with_error() -> None: assert any("Yuck" in message for message in log_messages) +@pytest.mark.integration @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 From 12040647365548e616bcb5e4a63eb9d6df177219 Mon Sep 17 00:00:00 2001 From: Mike Knepper Date: Fri, 22 May 2026 08:23:04 -0500 Subject: [PATCH 16/17] Suppress instead of try/except Signed-off-by: Mike Knepper --- .../src/nemo_data_designer_plugin/sdk/errors.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) 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 60e02cfd26..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 @@ -3,6 +3,8 @@ from __future__ import annotations +from contextlib import suppress + import httpx from data_designer.errors import DataDesignerError @@ -46,10 +48,9 @@ def extract_http_error_info(exc: httpx.HTTPStatusError) -> tuple[int, str]: pass detail = response.text - try: + body = None + with suppress(Exception): body = response.json() - except Exception: - body = None if isinstance(body, dict) and isinstance(body.get("detail"), str): detail = body["detail"] From d3e60acc42cbdfa90de1806de93c9af5f2353f6c Mon Sep 17 00:00:00 2001 From: Mike Knepper Date: Fri, 22 May 2026 08:23:19 -0500 Subject: [PATCH 17/17] Mark integration tests at module level Signed-off-by: Mike Knepper --- .../tests/integration/test_job_sdk.py | 18 ++---------------- .../tests/integration/test_job_task.py | 6 ++---- .../tests/integration/test_personas_cli.py | 10 ++-------- .../integration/test_preview_local_cli.py | 6 ++---- .../integration/test_preview_remote_sdk.py | 9 ++------- .../test_remote_validation_errors.py | 11 ++--------- .../test_sdk_get_model_providers.py | 4 ++-- 7 files changed, 14 insertions(+), 50 deletions(-) diff --git a/plugins/nemo-data-designer/tests/integration/test_job_sdk.py b/plugins/nemo-data-designer/tests/integration/test_job_sdk.py index b69420d8c6..ca74a22691 100644 --- a/plugins/nemo-data-designer/tests/integration/test_job_sdk.py +++ b/plugins/nemo-data-designer/tests/integration/test_job_sdk.py @@ -36,6 +36,8 @@ 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" @@ -97,7 +99,6 @@ def _no_pause() -> Generator[None]: # --------------------------------------------------------------------------- -@pytest.mark.integration @pytest.mark.asyncio async def test_get_job_resource_returns_job_for_real_job() -> None: async with _completed_job() as ctx: @@ -110,7 +111,6 @@ async def test_get_job_resource_returns_job_for_real_job() -> None: assert job["name"] == _JOB_NAME -@pytest.mark.integration @pytest.mark.asyncio async def test_get_job_resource_async_returns_job_for_real_job() -> None: async with _completed_job() as ctx: @@ -128,7 +128,6 @@ async def test_get_job_resource_async_returns_job_for_real_job() -> None: # --------------------------------------------------------------------------- -@pytest.mark.integration @pytest.mark.asyncio async def test_check_if_complete_returns_true_for_completed_status() -> None: async with _pending_job() as ctx: @@ -137,7 +136,6 @@ async def test_check_if_complete_returns_true_for_completed_status() -> None: assert job_resource.check_if_complete() is True -@pytest.mark.integration @pytest.mark.asyncio @pytest.mark.parametrize( ("simulated_status", "expected_log_fragment"), @@ -168,7 +166,6 @@ async def test_check_if_complete_returns_false_with_friendly_message_for_non_com ) -@pytest.mark.integration @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: @@ -185,7 +182,6 @@ async def test_check_if_complete_raises_when_requested(simulated_status: str) -> # --------------------------------------------------------------------------- -@pytest.mark.integration @pytest.mark.asyncio async def test_wait_until_done_logs_success_when_status_completes(caplog: pytest.LogCaptureFixture) -> None: async with _completed_job() as ctx: @@ -196,7 +192,6 @@ async def test_wait_until_done_logs_success_when_status_completes(caplog: pytest assert any("completed successfully" in record.message for record in caplog.records) -@pytest.mark.integration @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: @@ -208,7 +203,6 @@ async def test_wait_until_done_async_logs_success_when_status_completes(caplog: assert any("completed successfully" in record.message for record in caplog.records) -@pytest.mark.integration @pytest.mark.asyncio async def test_wait_until_done_logs_terminal_failure_for_cancelled_status( caplog: pytest.LogCaptureFixture, @@ -239,7 +233,6 @@ async def test_wait_until_done_logs_terminal_failure_for_cancelled_status( # --------------------------------------------------------------------------- -@pytest.mark.integration @pytest.mark.asyncio async def test_download_artifacts_extracts_dataset_and_loads_analysis(tmp_path: Path) -> None: async with _completed_job() as ctx: @@ -259,7 +252,6 @@ async def test_download_artifacts_extracts_dataset_and_loads_analysis(tmp_path: assert analysis.num_records == 3 -@pytest.mark.integration @pytest.mark.asyncio async def test_download_artifacts_async_extracts_dataset_and_loads_analysis(tmp_path: Path) -> None: async with _completed_job() as ctx: @@ -272,7 +264,6 @@ async def test_download_artifacts_async_extracts_dataset_and_loads_analysis(tmp_ assert results.load_analysis().num_records == 3 -@pytest.mark.integration @pytest.mark.asyncio async def test_load_processor_dataset_raises_for_unknown_processor(tmp_path: Path) -> None: async with _completed_job() as ctx: @@ -289,7 +280,6 @@ async def test_load_processor_dataset_raises_for_unknown_processor(tmp_path: Pat # --------------------------------------------------------------------------- -@pytest.mark.integration @pytest.mark.asyncio async def test_load_analysis_returns_profiler_results_for_completed_status() -> None: async with _completed_job() as ctx: @@ -301,7 +291,6 @@ async def test_load_analysis_returns_profiler_results_for_completed_status() -> assert analysis.num_records == 3 -@pytest.mark.integration @pytest.mark.asyncio async def test_load_analysis_raises_when_status_is_unknown() -> None: async with _pending_job() as ctx: @@ -312,7 +301,6 @@ async def test_load_analysis_raises_when_status_is_unknown() -> None: job_resource.load_analysis() -@pytest.mark.integration @pytest.mark.asyncio async def test_load_analysis_when_active_uses_completed_result_if_available( caplog: pytest.LogCaptureFixture, @@ -330,7 +318,6 @@ async def test_load_analysis_when_active_uses_completed_result_if_available( assert any("still cooking" in record.message.lower() for record in caplog.records) -@pytest.mark.integration @pytest.mark.asyncio async def test_load_analysis_when_terminally_incomplete_warns_and_returns_partial( caplog: pytest.LogCaptureFixture, @@ -345,7 +332,6 @@ async def test_load_analysis_when_terminally_incomplete_warns_and_returns_partia assert any("error" in record.message and "analysis" in record.message for record in caplog.records) -@pytest.mark.integration @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 diff --git a/plugins/nemo-data-designer/tests/integration/test_job_task.py b/plugins/nemo-data-designer/tests/integration/test_job_task.py index 7aae501a1c..68a6f6fee6 100644 --- a/plugins/nemo-data-designer/tests/integration/test_job_task.py +++ b/plugins/nemo-data-designer/tests/integration/test_job_task.py @@ -35,6 +35,8 @@ 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. @@ -63,7 +65,6 @@ def _failing_result_manager() -> Generator[None]: yield -@pytest.mark.integration @pytest.mark.asyncio async def test_task(tmp_path: Path) -> None: test_value = "test-value" @@ -99,7 +100,6 @@ async def test_task(tmp_path: Path) -> None: 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" @@ -137,7 +137,6 @@ async def test_save_partial_dataset_on_failure(_failing_result_manager: None, tm # 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.integration @pytest.mark.asyncio async def test_exiting_with_error() -> None: builder = dd.DataDesignerConfigBuilder(model_configs=[u.make_model_config()]) @@ -160,7 +159,6 @@ async def test_exiting_with_error() -> None: assert any("Yuck" in message for message in log_messages) -@pytest.mark.integration @pytest.mark.asyncio async def test_seed_dataset(tmp_path: Path) -> None: builder = dd.DataDesignerConfigBuilder(model_configs=[u.make_model_config()]) diff --git a/plugins/nemo-data-designer/tests/integration/test_personas_cli.py b/plugins/nemo-data-designer/tests/integration/test_personas_cli.py index 535a012a3c..8331f9242a 100644 --- a/plugins/nemo-data-designer/tests/integration/test_personas_cli.py +++ b/plugins/nemo-data-designer/tests/integration/test_personas_cli.py @@ -11,6 +11,8 @@ from nemo_platform import NeMoPlatform from nemo_platform.types.files import NGCStorageConfig +pytestmark = pytest.mark.integration + @pytest.fixture def mock_ngc_client() -> Generator[dict[str, Mock]]: @@ -44,7 +46,6 @@ def cli_sdk(monkeypatch: pytest.MonkeyPatch, sdk: NeMoPlatform) -> NeMoPlatform: return sdk -@pytest.mark.integration def test_personas_download_is_wired_properly() -> None: result = u.invoke_cli(["personas", "download", "--help"]) @@ -53,7 +54,6 @@ def test_personas_download_is_wired_properly() -> None: assert "data-designer download personas" not in result.output -@pytest.mark.integration def test_make_fileset_creates_requested_locale_with_existing_secret(cli_sdk: NeMoPlatform) -> None: result = u.invoke_cli( [ @@ -75,7 +75,6 @@ def test_make_fileset_creates_requested_locale_with_existing_secret(cli_sdk: NeM assert fileset.storage.api_key_secret == "system/ngc-api-key" -@pytest.mark.integration def test_make_fileset_creates_secret_from_env_then_fileset( monkeypatch: pytest.MonkeyPatch, cli_sdk: NeMoPlatform ) -> None: @@ -103,7 +102,6 @@ def test_make_fileset_creates_secret_from_env_then_fileset( assert fileset.storage.api_key_secret == "system/my-ngc-key" -@pytest.mark.integration def test_make_fileset_missing_env_var() -> None: result = u.invoke_cli( [ @@ -123,7 +121,6 @@ def test_make_fileset_missing_env_var() -> None: assert "not set or is empty" in result.output -@pytest.mark.integration def test_make_fileset_unknown_locale() -> None: result = u.invoke_cli( [ @@ -141,7 +138,6 @@ def test_make_fileset_unknown_locale() -> None: assert "de_DE" in result.output -@pytest.mark.integration def test_make_fileset_bare_secret_name() -> None: result = u.invoke_cli( [ @@ -158,7 +154,6 @@ def test_make_fileset_bare_secret_name() -> None: assert "WORKSPACE/NAME" in result.output -@pytest.mark.integration def test_make_fileset_create_secret_conflict_does_not_create_fileset( monkeypatch: pytest.MonkeyPatch, cli_sdk: NeMoPlatform ) -> None: @@ -184,7 +179,6 @@ def test_make_fileset_create_secret_conflict_does_not_create_fileset( assert filesets.data == [] -@pytest.mark.integration def test_make_fileset_create_secret_internal_error_surfaces_clearly( monkeypatch: pytest.MonkeyPatch, cli_sdk: NeMoPlatform ) -> None: diff --git a/plugins/nemo-data-designer/tests/integration/test_preview_local_cli.py b/plugins/nemo-data-designer/tests/integration/test_preview_local_cli.py index 8dbb773fd2..ac4cc78d75 100644 --- a/plugins/nemo-data-designer/tests/integration/test_preview_local_cli.py +++ b/plugins/nemo-data-designer/tests/integration/test_preview_local_cli.py @@ -9,8 +9,9 @@ from data_designer.cli.utils.sample_records_pager import PAGER_FILENAME from data_designer.config.analysis.dataset_profiler import DatasetProfilerResults +pytestmark = pytest.mark.integration + -@pytest.mark.integration def test_preview_run_saves_expected_artifacts(tmp_path: Path) -> None: config_path = _write_sampler_config(tmp_path) artifact_path = tmp_path / "preview-artifacts" @@ -40,7 +41,6 @@ def test_preview_run_saves_expected_artifacts(tmp_path: Path) -> None: assert (results_dir / "sample_records" / PAGER_FILENAME).exists() -@pytest.mark.integration 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) @@ -80,7 +80,6 @@ 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(tmp_path: Path) -> None: config_path = u.write_config_file( tmp_path, @@ -113,7 +112,6 @@ 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(tmp_path: Path) -> None: config_path = _write_sampler_config(tmp_path) diff --git a/plugins/nemo-data-designer/tests/integration/test_preview_remote_sdk.py b/plugins/nemo-data-designer/tests/integration/test_preview_remote_sdk.py index c535e740b4..c1d7fbaef5 100644 --- a/plugins/nemo-data-designer/tests/integration/test_preview_remote_sdk.py +++ b/plugins/nemo-data-designer/tests/integration/test_preview_remote_sdk.py @@ -13,8 +13,9 @@ from nemo_data_designer_plugin.config import get_config 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,7 +67,6 @@ def test_happy_path_preview() -> None: assert_message_with(log_messages, fuzzy="Preview generation in progress") -@pytest.mark.integration def test_hf_seed_dataset() -> None: builder = dd.DataDesignerConfigBuilder(model_configs=[u.make_model_config()]) builder.with_seed_dataset( @@ -87,7 +86,6 @@ def test_hf_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,7 +171,6 @@ def test_preview_with_schema_transform_processor() -> None: assert "messages" in processor_records[0] -@pytest.mark.integration 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 diff --git a/plugins/nemo-data-designer/tests/integration/test_remote_validation_errors.py b/plugins/nemo-data-designer/tests/integration/test_remote_validation_errors.py index 7c5f42c3a4..1b6cd1a059 100644 --- a/plugins/nemo-data-designer/tests/integration/test_remote_validation_errors.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, @@ -40,7 +42,6 @@ def _builder_with_llm_column(model_config: dd.ModelConfig) -> dd.DataDesignerCon return builder -@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) @@ -51,7 +52,6 @@ def test_unknown_provider_in_request() -> None: _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") @@ -62,7 +62,6 @@ def test_model_config_without_explicit_provider_is_rejected() -> None: _assert_error(dd_client, builder, ["does not have an explicit provider defined", alias]) -@pytest.mark.integration def test_malformed_provider_reference_is_rejected() -> None: alias = "too-many-slashes" malformed_provider_name = "foo/bar/baz" @@ -74,7 +73,6 @@ def test_malformed_provider_reference_is_rejected() -> None: _assert_error(dd_client, builder, ["Malformed model provider", alias, malformed_provider_name]) -@pytest.mark.integration 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) @@ -88,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" @@ -113,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) @@ -132,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()]) @@ -149,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( @@ -166,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 index 15da7f64ed..703917888a 100644 --- 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 @@ -5,8 +5,9 @@ import nemo_data_designer_plugin.testing.utils as u import pytest +pytestmark = pytest.mark.integration + -@pytest.mark.integration def test_get_default_model_providers_returns_registered_providers() -> None: """The SDK exposes IGW-registered providers as Data Designer ModelProviders.""" @@ -25,7 +26,6 @@ def test_get_default_model_providers_returns_registered_providers() -> None: assert provider.endpoint, f"Provider {provider.name!r} has no endpoint" -@pytest.mark.integration 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)."""