litellm/tests/unit/proxy/client/test_credentials.py

import time
import pytest
import requests



import responses

from litellm.proxy.client.credentials import CredentialsManagementClient
from litellm.proxy.client.exceptions import UnauthorizedError
from litellm.proxy.credential_endpoints.endpoints import CredentialHelperUtils
from litellm.types.utils import CredentialItem


@pytest.fixture
def base_url():
    return "http://localhost:8000"


@pytest.fixture
def api_key():
    return "test-api-key"


@pytest.fixture
def client(base_url, api_key):
    return CredentialsManagementClient(base_url=base_url, api_key=api_key)


def test_client_initialization(base_url, api_key):
    """Test that the CredentialsManagementClient is properly initialized"""
    client = CredentialsManagementClient(base_url=base_url, api_key=api_key)

    assert client._base_url == base_url
    assert client._api_key == api_key


def test_client_initialization_strips_trailing_slash():
    """Test that the client properly strips trailing slashes from base_url during initialization"""
    base_url = "http://localhost:8000/////"
    client = CredentialsManagementClient(base_url=base_url)

    assert client._base_url == "http://localhost:8000"


def test_client_without_api_key(base_url):
    """Test that the client works without an API key"""
    client = CredentialsManagementClient(base_url=base_url)

    assert client._api_key is None


def test_list_request(client, base_url, api_key):
    """Test list request"""
    request = client.list(return_request=True)

    assert request.method == "GET"
    assert request.url == f"{base_url}/credentials"
    assert request.headers["Content-Type"] == "application/json"
    assert request.headers["Authorization"] == f"Bearer {api_key}"


@responses.activate
def test_list_mock_response(client):
    """Test list with a mocked successful response"""
    mock_response = {
        "credentials": [
            {
                "credential_name": "azure1",
                "credential_info": {"api_type": "azure"},
                "model_id": "gpt-4",
            },
            {
                "credential_name": "anthropic1",
                "credential_info": {"api_type": "anthropic"},
                "model_id": "claude-2",
            },
        ]
    }

    responses.add(
        responses.GET,
        f"{client._base_url}/credentials",
        json=mock_response,
        status=200,
    )

    response = client.list()
    assert response == mock_response


@responses.activate
def test_list_unauthorized_error(client):
    """Test that list raises UnauthorizedError for 401 responses"""
    responses.add(
        responses.GET,
        f"{client._base_url}/credentials",
        status=401,
        json={"error": "Unauthorized"},
    )

    with pytest.raises(UnauthorizedError):
        client.list()


def test_create_request(client, base_url, api_key):
    """Test create request"""
    request = client.create(
        credential_name="azure1",
        credential_info={"api_type": "azure"},
        credential_values={
            "api_key": "sk-123",
            "api_base": "https://example.azure.openai.com",
        },
        return_request=True,
    )

    assert request.method == "POST"
    assert request.url == f"{base_url}/credentials"
    assert request.headers["Content-Type"] == "application/json"
    assert request.headers["Authorization"] == f"Bearer {api_key}"
    assert request.json == {
        "credential_name": "azure1",
        "credential_info": {"api_type": "azure"},
        "credential_values": {
            "api_key": "sk-123",
            "api_base": "https://example.azure.openai.com",
        },
    }


def test_create_request_includes_the_display_name_only_when_given(client):
    labeled = client.create("azure1", {}, {"api_key": "sk-123"}, return_request=True, display_name="Azure EU")
    unlabeled = client.create("azure1", {}, {"api_key": "sk-123"}, return_request=True)

    assert labeled.json["display_name"] == "Azure EU"
    assert "display_name" not in unlabeled.json


@pytest.mark.parametrize("display_name", ["Azure EU", None])
def test_update_display_name_request(client, base_url, api_key, display_name):
    request = client.update_display_name("azure1", display_name, return_request=True)

    assert request.method == "PATCH"
    assert request.url == f"{base_url}/credentials/azure1"
    assert request.headers["Authorization"] == f"Bearer {api_key}"
    assert request.json == {"display_name": display_name, "credential_info": {}}


@pytest.mark.parametrize(("credential_name", "path"), [("team?prod", "team%3Fprod"), ("a/b#c", "a%2Fb%23c")])
def test_update_display_name_targets_the_whole_credential_name(client, base_url, credential_name, path):
    request = client.update_display_name(credential_name, "Label", return_request=True)

    assert request.prepare().url == f"{base_url}/credentials/{path}"


@responses.activate
def test_create_mock_response(client):
    """Test create with a mocked successful response"""
    mock_response = {
        "credential_name": "azure1",
        "credential_info": {"api_type": "azure"},
        "status": "success",
    }

    responses.add(
        responses.POST,
        f"{client._base_url}/credentials",
        json=mock_response,
        status=200,
    )

    response = client.create(
        credential_name="azure1",
        credential_info={"api_type": "azure"},
        credential_values={
            "api_key": "sk-123",
            "api_base": "https://example.azure.openai.com",
        },
    )
    assert response == mock_response


@responses.activate
def test_create_unauthorized_error(client):
    """Test that create raises UnauthorizedError for 401 responses"""
    responses.add(
        responses.POST,
        f"{client._base_url}/credentials",
        status=401,
        json={"error": "Unauthorized"},
    )

    with pytest.raises(UnauthorizedError):
        client.create(
            credential_name="azure1",
            credential_info={"api_type": "azure"},
            credential_values={"api_key": "sk-123"},
        )


def test_delete_request(client, base_url, api_key):
    """Test delete request"""
    request = client.delete(credential_name="azure1", return_request=True)

    assert request.method == "DELETE"
    assert request.url == f"{base_url}/credentials/azure1"
    assert request.headers["Content-Type"] == "application/json"
    assert request.headers["Authorization"] == f"Bearer {api_key}"


@responses.activate
def test_delete_mock_response(client):
    """Test delete with a mocked successful response"""
    mock_response = {"credential_name": "azure1", "status": "deleted"}

    responses.add(
        responses.DELETE,
        f"{client._base_url}/credentials/azure1",
        json=mock_response,
        status=200,
    )

    response = client.delete(credential_name="azure1")
    assert response == mock_response


@responses.activate
def test_delete_unauthorized_error(client):
    """Test that delete raises UnauthorizedError for 401 responses"""
    responses.add(
        responses.DELETE,
        f"{client._base_url}/credentials/azure1",
        status=401,
        json={"error": "Unauthorized"},
    )

    with pytest.raises(UnauthorizedError):
        client.delete(credential_name="azure1")


def test_get_request(client, base_url, api_key):
    """Test get request"""
    request = client.get(credential_name="azure1", return_request=True)

    assert request.method == "GET"
    assert request.url == f"{base_url}/credentials/by_name/azure1"
    assert request.headers["Content-Type"] == "application/json"
    assert request.headers["Authorization"] == f"Bearer {api_key}"


@responses.activate
def test_get_mock_response(client):
    """Test get with a mocked successful response"""
    mock_response = {
        "credential_name": "azure1",
        "credential_info": {"api_type": "azure"},
        "credential_values": {
            "api_key": "sk-123",
            "api_base": "https://example.azure.openai.com",
        },
        "status": "active",
    }

    responses.add(
        responses.GET,
        f"{client._base_url}/credentials/by_name/azure1",
        json=mock_response,
        status=200,
    )

    response = client.get(credential_name="azure1")
    assert response == mock_response


@responses.activate
def test_get_unauthorized_error(client):
    """Test that get raises UnauthorizedError for 401 responses"""
    responses.add(
        responses.GET,
        f"{client._base_url}/credentials/by_name/azure1",
        status=401,
        json={"error": "Unauthorized"},
    )

    with pytest.raises(UnauthorizedError):
        client.get(credential_name="azure1")


def test_encrypt_credential_values_does_not_mutate_original(monkeypatch):
    """Ensure encrypt_credential_values returns a new encrypted object"""
    monkeypatch.setenv("LITELLM_SALT_KEY", "test-key")
    credential = CredentialItem(
        credential_name="azure1",
        credential_values={"api_key": "sk-123"},
        credential_info={"api_type": "azure"},
    )

    encrypted = CredentialHelperUtils.encrypt_credential_values(credential)

    assert encrypted.credential_values["api_key"] != "sk-123"
    assert credential.credential_values["api_key"] == "sk-123"
    assert encrypted.credential_name == credential.credential_name


def test_list_gives_up_at_the_timeout_instead_of_hanging(hanging_server):
    """
    A proxy that accepts the connection but never answers used to pin the caller's
    process forever, since the request carried no timeout at all.
    """
    client = CredentialsManagementClient(base_url=hanging_server, api_key="sk-test", timeout=1)

    started = time.monotonic()
    with pytest.raises(requests.exceptions.Timeout):
        client.list()

    assert time.monotonic() - started < 10