DataDesigner/packages/data-designer-engine/tests/engine/models/conftest.py

# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

from __future__ import annotations

from pathlib import Path
from unittest.mock import MagicMock

import pytest

from data_designer.config.models import (
    ChatCompletionInferenceParams,
    EmbeddingInferenceParams,
    ModelConfig,
)
from data_designer.engine.model_provider import ModelProvider, ModelProviderRegistry
from data_designer.engine.models.clients.base import ModelClient
from data_designer.engine.models.factory import create_model_registry
from data_designer.engine.models.registry import ModelRegistry
from data_designer.engine.secret_resolver import SecretsFileResolver
from data_designer.engine.testing import StubMCPFacade, StubMCPRegistry


@pytest.fixture
def stub_secrets_resolver() -> SecretsFileResolver:
    module_path = Path(__file__).parent
    return SecretsFileResolver(module_path / "stub_secrets.json")


@pytest.fixture
def stub_model_provider_registry() -> ModelProviderRegistry:
    return ModelProviderRegistry(
        providers=[
            ModelProvider(
                name="stub-model-provider",
                endpoint="https://api.example.com/v1",
                provider_type="openai",
                api_key="STUB_API_KEY",
            )
        ]
    )


@pytest.fixture
def stub_model_configs() -> list[ModelConfig]:
    return [
        ModelConfig(
            alias="stub-text",
            model="stub-model-text",
            provider="stub-model-provider",
            inference_parameters=ChatCompletionInferenceParams(
                temperature=0.80, top_p=0.95, max_tokens=100, max_parallel_requests=10, timeout=100
            ),
        ),
        ModelConfig(
            alias="stub-reasoning",
            model="stub-model-reasoning",
            provider="stub-model-provider",
            inference_parameters=ChatCompletionInferenceParams(
                temperature=0.80, top_p=0.95, max_tokens=100, max_parallel_requests=10, timeout=100
            ),
        ),
        ModelConfig(
            alias="stub-embedding",
            model="stub-model-embedding",
            provider="stub-model-provider",
            inference_parameters=EmbeddingInferenceParams(
                dimensions=100,
            ),
        ),
    ]


@pytest.fixture
def stub_model_registry(
    stub_model_configs: list[ModelConfig],
    stub_secrets_resolver: SecretsFileResolver,
    stub_model_provider_registry: ModelProviderRegistry,
) -> ModelRegistry:
    return create_model_registry(
        model_configs=stub_model_configs,
        secret_resolver=stub_secrets_resolver,
        model_provider_registry=stub_model_provider_registry,
    )


@pytest.fixture
def stub_model_client() -> MagicMock:
    """Mock ModelClient for testing ModelFacade without a real LiteLLM router."""
    return MagicMock(spec=ModelClient)


@pytest.fixture
def stub_mcp_facade_for_model() -> StubMCPFacade:
    """Default stub MCP facade with max_tool_call_turns=3."""
    return StubMCPFacade()


@pytest.fixture
def stub_mcp_registry_for_model(stub_mcp_facade_for_model: StubMCPFacade) -> StubMCPRegistry:
    """Default stub MCP registry."""
    return StubMCPRegistry(stub_mcp_facade_for_model)