DataDesigner/packages/data-designer/tests/cli/commands/test_create_command.py at 56fbe31f7c320cee311b14fd66a303db6940d881 · NVIDIA-NeMo/DataDesigner · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

from __future__ import annotations

from unittest.mock import MagicMock, patch

from data_designer.cli.commands.create import create_command

# ---------------------------------------------------------------------------
# create_command delegation tests
# ---------------------------------------------------------------------------


@patch("data_designer.cli.commands.create.GenerationController")
def test_create_command_delegates_to_controller(mock_ctrl_cls: MagicMock) -> None:
    """Test create_command delegates to GenerationController.run_create."""
    mock_ctrl = MagicMock()
    mock_ctrl_cls.return_value = mock_ctrl

    create_command(config_source="config.yaml", num_records=10, dataset_name="dataset", artifact_path=None, output_format=None)

    mock_ctrl_cls.assert_called_once()
    mock_ctrl.run_create.assert_called_once_with(
        config_source="config.yaml",
        num_records=10,
        dataset_name="dataset",
        artifact_path=None,
        output_format=None,
    )


@patch("data_designer.cli.commands.create.GenerationController")
def test_create_command_passes_custom_options(mock_ctrl_cls: MagicMock) -> None:
    """Test create_command passes custom options to the controller."""
    mock_ctrl = MagicMock()
    mock_ctrl_cls.return_value = mock_ctrl

    create_command(
        config_source="config.py",
        num_records=100,
        dataset_name="my_data",
        artifact_path="/custom/output",
        output_format=None,
    )

    mock_ctrl.run_create.assert_called_once_with(
        config_source="config.py",
        num_records=100,
        dataset_name="my_data",
        artifact_path="/custom/output",
        output_format=None,
    )


@patch("data_designer.cli.commands.create.GenerationController")
def test_create_command_default_artifact_path_is_none(mock_ctrl_cls: MagicMock) -> None:
    """Test create_command passes artifact_path=None when not specified."""
    mock_ctrl = MagicMock()
    mock_ctrl_cls.return_value = mock_ctrl

    create_command(config_source="config.yaml", num_records=5, dataset_name="ds", artifact_path=None, output_format=None)

    mock_ctrl.run_create.assert_called_once_with(
        config_source="config.yaml",
        num_records=5,
        dataset_name="ds",
        artifact_path=None,
        output_format=None,
    )


@patch("data_designer.cli.commands.create.GenerationController")
def test_create_command_passes_output_format(mock_ctrl_cls: MagicMock) -> None:
    """Test create_command forwards --output-format to the controller."""
    mock_ctrl = MagicMock()
    mock_ctrl_cls.return_value = mock_ctrl

    create_command(
        config_source="config.yaml",
        num_records=10,
        dataset_name="dataset",
        artifact_path=None,
        output_format="jsonl",
    )

    mock_ctrl.run_create.assert_called_once_with(
        config_source="config.yaml",
        num_records=10,
        dataset_name="dataset",
        artifact_path=None,
        output_format="jsonl",
    )