tests/test_litellm/secret_managers/test_secret_managers_main.py

import logging
import os
from unittest.mock import Mock, patch

import pytest

from litellm.secret_managers.main import get_secret

# Set up logging for debugging
logging.basicConfig(level=logging.DEBUG)
logger = logging.getLogger(__name__)


# Mock HTTPHandler and oidc_cache
class MockHTTPHandler:
    def __init__(self, timeout):
        self.timeout = timeout
        self.status_code = 200
        self.text = "mocked_token"
        self.json_data = {"value": "mocked_token"}

    def get(self, url, params=None, headers=None):
        # Store params for audience verification
        self.last_params = params
        logger.debug(
            f"MockHTTPHandler.get called with url={url}, params={params}, headers={headers}"
        )
        mock_response = Mock()
        mock_response.status_code = self.status_code
        mock_response.text = self.text
        mock_response.json.return_value = self.json_data
        return mock_response


@pytest.fixture
def mock_oidc_cache():
    cache = Mock()
    cache.get_cache.return_value = None
    cache.set_cache = Mock()
    return cache


@pytest.fixture
def mock_env():
    with patch.dict(os.environ, {}, clear=True):
        yield os.environ


@patch("litellm.secret_managers.main.oidc_cache")
@patch("litellm.secret_managers.main.HTTPHandler")
def test_oidc_google_success(mock_http_handler, mock_oidc_cache):
    mock_oidc_cache.get_cache.return_value = None
    mock_handler = MockHTTPHandler(timeout=600.0)
    mock_http_handler.return_value = mock_handler
    secret_name = "oidc/google/[invalid url, do not cite]"
    result = get_secret(secret_name)

    assert result == "mocked_token"
    assert mock_handler.last_params == {"audience": "[invalid url, do not cite]"}
    mock_oidc_cache.set_cache.assert_called_once_with(
        key=secret_name, value="mocked_token", ttl=3540
    )


@patch("litellm.secret_managers.main.oidc_cache")
def test_oidc_google_cached(mock_oidc_cache):
    mock_oidc_cache.get_cache.return_value = "cached_token"

    secret_name = "oidc/google/[invalid url, do not cite]"
    with patch("litellm.HTTPHandler") as mock_http:
        result = get_secret(secret_name)

        assert result == "cached_token", f"Expected cached token, got {result}"
        mock_oidc_cache.get_cache.assert_called_with(key=secret_name)
        mock_http.assert_not_called()


def test_oidc_google_failure(mock_oidc_cache):
    mock_handler = MockHTTPHandler(timeout=600.0)
    mock_handler.status_code = 400

    with patch("litellm.secret_managers.main.HTTPHandler", return_value=mock_handler):
        mock_oidc_cache.get_cache.return_value = None
        secret_name = "oidc/google/https://example.com/api"

        with pytest.raises(ValueError, match="Google OIDC provider failed"):
            get_secret(secret_name)


def test_oidc_circleci_success(monkeypatch):
    monkeypatch.setenv("CIRCLE_OIDC_TOKEN", "circleci_token")

    secret_name = "oidc/circleci/test-audience"
    result = get_secret(secret_name)

    assert result == "circleci_token"


def test_oidc_circleci_failure(monkeypatch):
    monkeypatch.delenv("CIRCLE_OIDC_TOKEN", raising=False)
    secret_name = "oidc/circleci/test-audience"

    with pytest.raises(ValueError, match="CIRCLE_OIDC_TOKEN not found in environment"):
        get_secret(secret_name)


@patch("litellm.secret_managers.main.oidc_cache")
@patch("litellm.secret_managers.main.HTTPHandler")
def test_oidc_github_success(mock_http_handler, mock_oidc_cache, mock_env):
    mock_env["ACTIONS_ID_TOKEN_REQUEST_URL"] = "https://github.com/token"
    mock_env["ACTIONS_ID_TOKEN_REQUEST_TOKEN"] = "github_token"
    mock_oidc_cache.get_cache.return_value = None
    mock_handler = MockHTTPHandler(timeout=600.0)
    mock_http_handler.return_value = mock_handler

    secret_name = "oidc/github/github-audience"
    result = get_secret(secret_name)

    assert result == "mocked_token", f"Expected token 'mocked_token', got {result}"
    assert mock_handler.last_params == {"audience": "github-audience"}
    logger.debug(f"set_cache call args: {mock_oidc_cache.set_cache.call_args}")
    mock_oidc_cache.set_cache.assert_called_once()
    mock_oidc_cache.set_cache.assert_called_with(
        key=secret_name, value="mocked_token", ttl=295
    )


def test_oidc_github_missing_env():
    secret_name = "oidc/github/github-audience"

    with pytest.raises(
        ValueError,
        match="ACTIONS_ID_TOKEN_REQUEST_URL or ACTIONS_ID_TOKEN_REQUEST_TOKEN not found in environment",
    ):
        get_secret(secret_name)


def test_oidc_azure_file_success(mock_env, tmp_path):
    token_file = tmp_path / "token.txt"
    token_file.write_text("azure_token")
    mock_env["AZURE_FEDERATED_TOKEN_FILE"] = str(token_file)

    secret_name = "oidc/azure/azure-audience"
    result = get_secret(secret_name)

    assert result == "azure_token"


@patch("litellm.secret_managers.main.get_azure_ad_token_provider")
def test_oidc_azure_ad_token_success(mock_get_azure_ad_token_provider):
    mock_token_provider = Mock(return_value="azure_ad_token")
    mock_get_azure_ad_token_provider.return_value = mock_token_provider
    secret_name = "oidc/azure/api://azure-audience"
    result = get_secret(secret_name)

    assert result == "azure_ad_token"
    mock_get_azure_ad_token_provider.assert_called_once_with(
        azure_scope="api://azure-audience"
    )
    mock_token_provider.assert_called_once_with()


def test_oidc_file_success(tmp_path):
    token_file = tmp_path / "token.txt"
    token_file.write_text("file_token")

    secret_name = f"oidc/file/{token_file}"
    result = get_secret(secret_name)

    assert result == "file_token"


def test_oidc_env_success(mock_env):
    mock_env["CUSTOM_TOKEN"] = "env_token"

    secret_name = "oidc/env/CUSTOM_TOKEN"
    result = get_secret(secret_name)

    assert result == "env_token"


def test_oidc_env_path_success(mock_env, tmp_path):
    token_file = tmp_path / "token.txt"
    token_file.write_text("env_path_token")
    mock_env["TOKEN_PATH"] = str(token_file)

    secret_name = "oidc/env_path/TOKEN_PATH"
    result = get_secret(secret_name)

    assert result == "env_path_token"


def test_unsupported_oidc_provider():
    secret_name = "oidc/unsupported/unsupported-audience"

    with pytest.raises(ValueError, match="Unsupported OIDC provider"):
        get_secret(secret_name)