new tests for cli components

Author: nabinchha

src/data_designer/cli/constants.py

@@ -9,3 +9,25 @@
 DEFAULT_CONFIG_DIR = Path.home() / ".data-designer"
 MODEL_CONFIGS_FILE_NAME = "model_configs.yaml"
 MODEL_PROVIDERS_FILE_NAME = "model_providers.yaml"
+
+# Predefined provider templates
+PREDEFINED_PROVIDERS = {
+    "nvidia": {
+        "name": "nvidia",
+        "endpoint": "https://integrate.api.nvidia.com/v1",
+        "provider_type": "openai",
+        "api_key": "NVIDIA_API_KEY",
+    },
+    "openai": {
+        "name": "openai",
+        "endpoint": "https://api.openai.com/v1",
+        "provider_type": "openai",
+        "api_key": "OPENAI_API_KEY",
+    },
+    "anthropic": {
+        "name": "anthropic",
+        "endpoint": "https://api.anthropic.com/v1",
+        "provider_type": "openai",
+        "api_key": "ANTHROPIC_API_KEY",
+    },
+}

src/data_designer/cli/forms/provider_builder.py

@@ -3,34 +3,13 @@
 
 from typing import Any
 
+from data_designer.cli.constants import PREDEFINED_PROVIDERS
 from data_designer.cli.forms.builder import FormBuilder
 from data_designer.cli.forms.field import SelectField, TextField
 from data_designer.cli.forms.form import Form
 from data_designer.cli.utils import validate_url
 from data_designer.engine.model_provider import ModelProvider
 
-# Predefined provider templates
-PREDEFINED_PROVIDERS = {
-    "nvidia": {
-        "name": "nvidia",
-        "endpoint": "https://integrate.api.nvidia.com/v1",
-        "provider_type": "openai",
-        "api_key": "NVIDIA_API_KEY",
-    },
-    "openai": {
-        "name": "openai",
-        "endpoint": "https://api.openai.com/v1",
-        "provider_type": "openai",
-        "api_key": "OPENAI_API_KEY",
-    },
-    "anthropic": {
-        "name": "anthropic",
-        "endpoint": "https://api.anthropic.com/v1",
-        "provider_type": "openai",
-        "api_key": "ANTHROPIC_API_KEY",
-    },
-}
-
 
 class ProviderFormBuilder(FormBuilder[ModelProvider]):
     """Builds interactive forms for provider configuration."""

src/data_designer/cli/repositories/model_repository.py

@@ -3,21 +3,18 @@
 
 from pathlib import Path
 
+from pydantic import BaseModel
+
 from data_designer.cli.constants import MODEL_CONFIGS_FILE_NAME
 from data_designer.cli.repositories.base import ConfigRepository
 from data_designer.cli.utils import load_config_file, save_config_file
 from data_designer.config.models import ModelConfig
 
 
-class ModelConfigRegistry:
+class ModelConfigRegistry(BaseModel):
     """Registry for model configurations."""
 
-    def __init__(self, model_configs: list[ModelConfig]):
-        self.model_configs = model_configs
-
-    def model_dump(self, **kwargs) -> dict:
-        """Dump to dictionary format."""
-        return {"model_configs": [mc.model_dump(**kwargs) for mc in self.model_configs]}
+    model_configs: list[ModelConfig]
 
 
 class ModelRepository(ConfigRepository[ModelConfigRegistry]):
@@ -35,10 +32,7 @@ def load(self) -> ModelConfigRegistry | None:
 
         try:
             config_dict = load_config_file(self.config_file)
-            if "model_configs" not in config_dict:
-                return None
-            model_configs = [ModelConfig.model_validate(mc) for mc in config_dict["model_configs"]]
-            return ModelConfigRegistry(model_configs)
+            return ModelConfigRegistry.model_validate(config_dict)
         except Exception:
             return None
 

src/data_designer/cli/repositories/provider_repository.py

@@ -3,10 +3,19 @@
 
 from pathlib import Path
 
+from pydantic import BaseModel
+
 from data_designer.cli.constants import MODEL_PROVIDERS_FILE_NAME
 from data_designer.cli.repositories.base import ConfigRepository
 from data_designer.cli.utils import load_config_file, save_config_file
-from data_designer.engine.model_provider import ModelProviderRegistry
+from data_designer.config.models import ModelProvider
+
+
+class ModelProviderRegistry(BaseModel):
+    """Registry for model provider configurations."""
+
+    providers: list[ModelProvider]
+    default: str | None = None
 
 
 class ProviderRepository(ConfigRepository[ModelProviderRegistry]):

src/data_designer/cli/services/model_service.py

@@ -29,7 +29,6 @@ def add(self, model: ModelConfig) -> None:
         """
         registry = self.repository.load() or ModelConfigRegistry(model_configs=[])
 
-        # Business rule: No duplicate aliases
         if any(m.alias == model.alias for m in registry.model_configs):
             raise ValueError(f"Model alias '{model.alias}' already exists")
 
@@ -54,12 +53,10 @@ def update(self, original_alias: str, updated_model: ModelConfig) -> None:
         if index is None:
             raise ValueError(f"Model '{original_alias}' not found")
 
-        # Business rule: Alias change must not conflict
         if updated_model.alias != original_alias:
             if any(m.alias == updated_model.alias for m in registry.model_configs):
                 raise ValueError(f"Model alias '{updated_model.alias}' already exists")
 
-        # Update
         registry.model_configs[index] = updated_model
         self.repository.save(registry)
 
@@ -76,10 +73,8 @@ def delete(self, alias: str) -> None:
         if not any(m.alias == alias for m in registry.model_configs):
             raise ValueError(f"Model '{alias}' not found")
 
-        # Remove model
         registry.model_configs = [m for m in registry.model_configs if m.alias != alias]
 
-        # Business rule: Delete file if no models left
         if registry.model_configs:
             self.repository.save(registry)
         else:

src/data_designer/cli/services/provider_service.py

@@ -1,8 +1,8 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-from data_designer.cli.repositories.provider_repository import ProviderRepository
-from data_designer.engine.model_provider import ModelProvider, ModelProviderRegistry
+from data_designer.cli.repositories.provider_repository import ModelProviderRegistry, ProviderRepository
+from data_designer.config.models import ModelProvider
 
 
 class ProviderService:
@@ -27,22 +27,12 @@ def add(self, provider: ModelProvider) -> None:
         Raises:
             ValueError: If provider name already exists
         """
-        registry = self.repository.load()
-
-        if registry:
-            # Business rule: No duplicate names
-            if any(p.name == provider.name for p in registry.providers):
-                raise ValueError(f"Provider '{provider.name}' already exists")
-
-            registry.providers.append(provider)
-        else:
-            # Create new registry with first provider
-            registry = ModelProviderRegistry(providers=[provider], default=provider.name)
+        registry = self.repository.load() or ModelProviderRegistry(providers=[], default=None)
 
-        # Business rule: First provider is default (for existing registries adding first provider)
-        if len(registry.providers) == 1 and registry.default is None:
-            registry.default = provider.name
+        if any(p.name == provider.name for p in registry.providers):
+            raise ValueError(f"Provider '{provider.name}' already exists")
 
+        registry.providers.append(provider)
         self.repository.save(registry)
 
     def update(self, original_name: str, updated_provider: ModelProvider) -> None:
@@ -63,15 +53,12 @@ def update(self, original_name: str, updated_provider: ModelProvider) -> None:
         if index is None:
             raise ValueError(f"Provider '{original_name}' not found")
 
-        # Business rule: Name change must not conflict
         if updated_provider.name != original_name:
             if any(p.name == updated_provider.name for p in registry.providers):
                 raise ValueError(f"Provider name '{updated_provider.name}' already exists")
 
-        # Update
         registry.providers[index] = updated_provider
 
-        # Business rule: Update default if name changed
         if registry.default == original_name and updated_provider.name != original_name:
             registry.default = updated_provider.name
 
@@ -90,14 +77,11 @@ def delete(self, name: str) -> None:
         if not any(p.name == name for p in registry.providers):
             raise ValueError(f"Provider '{name}' not found")
 
-        # Remove provider
         registry.providers = [p for p in registry.providers if p.name != name]
 
-        # Business rule: Update default if deleted
         if registry.default == name:
             registry.default = registry.providers[0].name if registry.providers else None
 
-        # Business rule: Delete file if no providers left
         if registry.providers:
             self.repository.save(registry)
         else:

tests/cli/conftest.py

@@ -0,0 +1,89 @@
+from pathlib import Path
+
+import pytest
+
+from data_designer.cli.repositories.model_repository import ModelConfigRegistry, ModelRepository
+from data_designer.cli.repositories.provider_repository import ModelProviderRegistry, ProviderRepository
+from data_designer.cli.services.model_service import ModelService
+from data_designer.cli.services.provider_service import ProviderService
+from data_designer.config.models import InferenceParameters, ModelConfig, ModelProvider
+
+
+@pytest.fixture
+def stub_model_configs() -> list[ModelConfig]:
+    return [
+        ModelConfig(
+            alias="test-alias-1",
+            model="test-model-1",
+            provider="test-provider-1",
+            inference_parameters=InferenceParameters(
+                temperature=0.7, top_p=0.9, max_tokens=2048, max_parallel_requests=4
+            ),
+        ),
+        ModelConfig(
+            alias="test-alias-2",
+            model="test-model-2",
+            provider="test-provider-1",
+            inference_parameters=InferenceParameters(
+                temperature=0.7, top_p=0.9, max_tokens=2048, max_parallel_requests=4
+            ),
+        ),
+    ]
+
+
+@pytest.fixture
+def stub_new_model_config() -> ModelConfig:
+    return ModelConfig(
+        alias="test-alias-3",
+        model="test-model-3",
+        provider="test-provider-1",
+        inference_parameters=InferenceParameters(
+            temperature=0.7,
+            top_p=0.9,
+            max_tokens=2048,
+            max_parallel_requests=4,
+            timeout=100,
+        ),
+    )
+
+
+@pytest.fixture
+def stub_model_providers() -> list[ModelProvider]:
+    return [
+        ModelProvider(
+            name="test-provider-1",
+            endpoint="https://api.example.com/v1",
+            provider_type="openai",
+            api_key="test-api-key",
+        ),
+        ModelProvider(
+            name="test-provider-2",
+            endpoint="https://api.example.com/v2",
+            provider_type="openai",
+            api_key="test-api-key-2",
+        ),
+    ]
+
+
+@pytest.fixture
+def stub_new_model_provider() -> ModelProvider:
+    return ModelProvider(
+        name="test-provider-3",
+        endpoint="https://api.example.com/v1",
+        provider_type="openai",
+        api_key="test-api-key-1",
+    )
+
+
+@pytest.fixture
+def stub_model_service(tmp_path: Path, stub_model_configs: list[ModelConfig]) -> ModelService:
+    repository = ModelRepository(tmp_path)
+    repository.save(ModelConfigRegistry(model_configs=stub_model_configs))
+    return ModelService(repository)
+
+
+@pytest.fixture
+def stub_provider_service(tmp_path: Path, stub_model_providers: list[ModelProvider]) -> ProviderService:
+    repository = ProviderRepository(tmp_path)
+    repository.save(ModelProviderRegistry(providers=stub_model_providers, default=stub_model_providers[0].name))
+    return ProviderService(repository)

tests/cli/repositories/test_model_repository.py

@@ -0,0 +1,31 @@
+from pathlib import Path
+
+from data_designer.cli.constants import MODEL_CONFIGS_FILE_NAME
+from data_designer.cli.repositories.model_repository import ModelConfigRegistry, ModelRepository
+from data_designer.cli.utils import save_config_file
+from data_designer.config.models import ModelConfig
+
+
+def test_config_file(tmp_path: Path):
+    repository = ModelRepository(tmp_path)
+    assert repository.config_file == tmp_path / MODEL_CONFIGS_FILE_NAME
+
+
+def test_load_does_not_exist(tmp_path: Path):
+    repository = ModelRepository(tmp_path)
+    assert repository.load() is None
+
+
+def test_load_exists(tmp_path: Path, stub_model_configs: list[ModelConfig]):
+    config_file = tmp_path / MODEL_CONFIGS_FILE_NAME
+    save_config_file(config_file, ModelConfigRegistry(model_configs=stub_model_configs).model_dump())
+    repository = ModelRepository(tmp_path)
+    assert repository.load() is not None
+    assert repository.load().model_configs == stub_model_configs
+
+
+def test_save(tmp_path: Path, stub_model_configs: list[ModelConfig]):
+    repository = ModelRepository(tmp_path)
+    repository.save(ModelConfigRegistry(model_configs=stub_model_configs))
+    assert repository.load() is not None
+    assert repository.load().model_configs == stub_model_configs

tests/cli/repositories/test_provider_repository.py

@@ -0,0 +1,34 @@
+from pathlib import Path
+
+from data_designer.cli.constants import MODEL_PROVIDERS_FILE_NAME
+from data_designer.cli.repositories.provider_repository import ModelProviderRegistry, ProviderRepository
+from data_designer.cli.utils import save_config_file
+from data_designer.config.models import ModelProvider
+
+
+def test_config_file(tmp_path: Path):
+    repository = ProviderRepository(tmp_path)
+    assert repository.config_file == tmp_path / MODEL_PROVIDERS_FILE_NAME
+
+
+def test_load_does_not_exist(tmp_path: Path):
+    repository = ProviderRepository(tmp_path)
+    assert repository.load() is None
+
+
+def test_load_exists(tmp_path: Path, stub_model_providers: list[ModelProvider]):
+    config_file_path = tmp_path / MODEL_PROVIDERS_FILE_NAME
+    save_config_file(
+        config_file_path,
+        ModelProviderRegistry(providers=stub_model_providers, default=stub_model_providers[0].name).model_dump(),
+    )
+    repository = ProviderRepository(tmp_path)
+    assert repository.load() is not None
+    assert repository.load().providers == stub_model_providers
+
+
+def test_save(tmp_path: Path, stub_model_providers: list[ModelProvider]):
+    repository = ProviderRepository(tmp_path)
+    repository.save(ModelProviderRegistry(providers=stub_model_providers, default=stub_model_providers[0].name))
+    assert repository.load() is not None
+    assert repository.load().providers == stub_model_providers