litellm/tests/unit/test_setup_wizard.py

"""Unit tests for litellm.setup_wizard — pure functions only, no network calls."""

from litellm.setup_wizard import SetupWizard, _yaml_escape

# ---------------------------------------------------------------------------
# _yaml_escape
# ---------------------------------------------------------------------------


def test_yaml_escape_plain():
    assert _yaml_escape("sk-abc123") == "sk-abc123"


def test_yaml_escape_double_quote():
    assert _yaml_escape('sk-ab"cd') == 'sk-ab\\"cd'


def test_yaml_escape_backslash():
    assert _yaml_escape("sk-ab\\cd") == "sk-ab\\\\cd"


def test_yaml_escape_combined():
    assert _yaml_escape('ab\\"cd') == 'ab\\\\\\"cd'


def test_yaml_escape_newline():
    assert _yaml_escape("sk-abc\ndef") == "sk-abc\\ndef"


def test_yaml_escape_carriage_return():
    assert _yaml_escape("sk-abc\rdef") == "sk-abc\\rdef"


def test_yaml_escape_tab():
    assert _yaml_escape("sk-abc\tdef") == "sk-abc\\tdef"


# ---------------------------------------------------------------------------
# SetupWizard._build_config
# ---------------------------------------------------------------------------

_OPENAI = {
    "id": "openai",
    "name": "OpenAI",
    "env_key": "OPENAI_API_KEY",
    "models": ["gpt-4o", "gpt-4o-mini"],
    "test_model": "gpt-4o-mini",
}

_ANTHROPIC = {
    "id": "anthropic",
    "name": "Anthropic",
    "env_key": "ANTHROPIC_API_KEY",
    "models": ["claude-opus-4-6"],
    "test_model": "claude-haiku-4-5-20251001",
}

_AZURE = {
    "id": "azure",
    "name": "Azure OpenAI",
    "env_key": "AZURE_AI_API_KEY",
    "models": [],
    "test_model": None,
    "needs_api_base": True,
    "api_base_hint": "https://<resource>.openai.azure.com/",
    "api_version": "2024-07-01-preview",
}

_OLLAMA = {
    "id": "ollama",
    "name": "Ollama",
    "env_key": None,
    "models": ["ollama/llama3.2"],
    "test_model": None,
    "api_base": "http://localhost:11434",
}


def test_build_config_basic_openai():
    config = SetupWizard._build_config(
        [_OPENAI],
        {"OPENAI_API_KEY": "sk-test"},
        "sk-master",
    )
    assert "model_list:" in config
    assert "model_name: gpt-4o" in config
    assert "model: gpt-4o" in config
    assert "api_key: os.environ/OPENAI_API_KEY" in config
    assert 'master_key: "sk-master"' in config


def test_build_config_skipped_provider_omitted():
    """Provider with no key in env_vars should not appear in model_list."""
    config = SetupWizard._build_config(
        [_OPENAI, _ANTHROPIC],
        {"ANTHROPIC_API_KEY": "sk-ant-test"},  # OpenAI key missing
        "sk-master",
    )
    assert "gpt-4o" not in config
    assert "claude-opus-4-6" in config


def test_build_config_env_vars_written_escaped():
    """API keys with special chars should be YAML-escaped."""
    config = SetupWizard._build_config(
        [_OPENAI],
        {"OPENAI_API_KEY": 'sk-ab"cd'},
        "sk-master",
    )
    assert 'OPENAI_API_KEY: "sk-ab\\"cd"' in config


def test_build_config_master_key_quoted():
    """master_key must be quoted in YAML to handle special characters."""
    config = SetupWizard._build_config(
        [_OPENAI],
        {"OPENAI_API_KEY": "sk-test"},
        'sk-master"special',
    )
    assert 'master_key: "sk-master\\"special"' in config


def test_build_config_does_not_mutate_env_vars():
    """_build_config must not modify the caller's env_vars dict."""
    env_vars = {
        "AZURE_AI_API_KEY": "az-key",
        "_LITELLM_AZURE_AI_API_BASE_AZURE": "https://my.azure.com",
        "_LITELLM_AZURE_DEPLOYMENT_AZURE": "my-deployment",
    }
    original_keys = set(env_vars.keys())
    SetupWizard._build_config([_AZURE], env_vars, "sk-master")
    assert set(env_vars.keys()) == original_keys


def test_build_config_azure_uses_deployment_name():
    env_vars = {
        "AZURE_AI_API_KEY": "az-key",
        "_LITELLM_AZURE_AI_API_BASE_AZURE": "https://my.azure.com",
        "_LITELLM_AZURE_DEPLOYMENT_AZURE": "my-gpt4o",
    }
    config = SetupWizard._build_config([_AZURE], env_vars, "sk-master")
    assert "model: azure/my-gpt4o" in config
    assert "model_name: azure-my-gpt4o" in config
    # api_base must be quoted to survive YAML special chars
    assert 'api_base: "https://my.azure.com"' in config


def test_build_config_azure_no_deployment_skipped():
    """Azure without a deployment name should emit nothing (not fallback to gpt-4o)."""
    env_vars = {"AZURE_AI_API_KEY": "az-key"}  # no deployment sentinel
    config = SetupWizard._build_config([_AZURE], env_vars, "sk-master")
    # No azure model entry should be emitted when deployment name is absent
    assert "model: azure/" not in config


def test_build_config_no_display_name_collision_openai_and_azure():
    """OpenAI gpt-4o and azure gpt-4o should get distinct model_name values."""
    env_vars = {
        "OPENAI_API_KEY": "sk-openai",
        "AZURE_AI_API_KEY": "az-key",
        "_LITELLM_AZURE_DEPLOYMENT_AZURE": "gpt-4o",
    }
    config = SetupWizard._build_config([_OPENAI, _AZURE], env_vars, "sk-master")
    assert "model_name: gpt-4o" in config  # OpenAI
    assert "model_name: azure-gpt-4o" in config  # Azure — qualified


def test_build_config_ollama_no_api_key_line():
    """Ollama has no env_key — config should not contain an api_key line for it."""
    config = SetupWizard._build_config([_OLLAMA], {}, "sk-master")
    assert "ollama/llama3.2" in config
    assert "api_key:" not in config


def test_build_config_master_key_in_general_settings():
    """master_key is written to general_settings."""
    config = SetupWizard._build_config([_OPENAI], {"OPENAI_API_KEY": "k"}, "sk-m")
    assert 'master_key: "sk-m"' in config


def test_build_config_internal_sentinel_keys_excluded():
    """_LITELLM_ prefixed sentinel keys must not appear in environment_variables."""
    env_vars = {
        "OPENAI_API_KEY": "sk-real",
        "_LITELLM_AZURE_AI_API_BASE_AZURE": "https://x.azure.com",
    }
    config = SetupWizard._build_config([_OPENAI], env_vars, "sk-master")
    assert "_LITELLM_" not in config