litellm-mirror/tests/litellm/secrets_managers/test_main_oidc.py

import pytest
import os
from unittest.mock import Mock, patch
from litellm.secret_managers.main import get_secret
import logging

# 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(mock_env):
    mock_env["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():
    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)