mirror of
https://github.com/NVIDIA-NeMo/DataDesigner
synced 2026-05-24 09:48:29 +00:00
8b79b21298
Preserves tree from previous docs-website head: 5e47d33ea8. This branch is a CI-managed publish artifact like gh-pages; source provenance is tracked in commit messages rather than Git ancestry.
57 lines
1.9 KiB
Python
57 lines
1.9 KiB
Python
# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
from __future__ import annotations
|
|
|
|
from unittest.mock import Mock
|
|
|
|
import pytest
|
|
|
|
import data_designer.lazy_heavy_imports as lazy
|
|
from data_designer.config.run_config import RunConfig
|
|
from data_designer.engine.models.facade import ModelFacade
|
|
from data_designer.engine.models.registry import ModelRegistry
|
|
from data_designer.engine.resources.person_reader import PersonReader
|
|
from data_designer.engine.resources.resource_provider import ResourceProvider
|
|
from data_designer.engine.storage.artifact_storage import ArtifactStorage
|
|
|
|
|
|
@pytest.fixture
|
|
def artifact_storage(tmp_path):
|
|
return ArtifactStorage(artifact_path=tmp_path)
|
|
|
|
|
|
@pytest.fixture
|
|
def stub_model_facade():
|
|
mock_facade = Mock(spec=ModelFacade)
|
|
mock_facade.model_alias = "test_model"
|
|
mock_facade.generate.return_value = ("Generated summary text", None)
|
|
return mock_facade
|
|
|
|
|
|
@pytest.fixture
|
|
def stub_resource_provider(tmp_path, stub_model_facade):
|
|
mock_provider = Mock(spec=ResourceProvider)
|
|
mock_model_registry = Mock(spec=ModelRegistry)
|
|
mock_model_registry.get_model.return_value = stub_model_facade
|
|
mock_model_registry.model_configs = {} # Add empty model_configs dict
|
|
mock_provider.model_registry = mock_model_registry
|
|
mock_provider.artifact_storage = ArtifactStorage(artifact_path=tmp_path)
|
|
mock_provider.person_reader = Mock(spec=PersonReader)
|
|
mock_provider.seed_reader = Mock()
|
|
mock_provider.seed_reader.get_column_names.return_value = []
|
|
mock_provider.run_config = RunConfig()
|
|
return mock_provider
|
|
|
|
|
|
@pytest.fixture
|
|
def stub_sample_dataframe():
|
|
return lazy.pd.DataFrame(
|
|
{
|
|
"col1": [1, 2, 3, 4],
|
|
"col2": ["a", "b", "c", "d"],
|
|
"col3": [True, False, True, False],
|
|
"category": ["A", "B", "A", "B"],
|
|
"other_col": [1, 2, 3, 4],
|
|
}
|
|
)
|