mirror of
https://github.com/tiennm99/litellm.git
synced 2026-07-19 10:19:07 +00:00
243 lines
9.5 KiB
Python
243 lines
9.5 KiB
Python
import os
|
|
import sys
|
|
from unittest.mock import MagicMock, call, patch
|
|
|
|
import pytest
|
|
|
|
sys.path.insert(
|
|
0, os.path.abspath("../../..")
|
|
) # Adds the parent directory to the system path
|
|
|
|
from litellm.llms.watsonx.common_utils import generate_iam_token
|
|
|
|
|
|
class TestGenerateIAMToken:
|
|
"""Tests for the generate_iam_token function, specifically testing API key fallback logic."""
|
|
|
|
@patch("litellm.llms.watsonx.common_utils.iam_token_cache")
|
|
@patch("litellm.llms.watsonx.common_utils.litellm.module_level_client")
|
|
@patch("litellm.llms.watsonx.common_utils.get_secret_str")
|
|
def test_generate_iam_token_with_watsonx_zenapikey(
|
|
self, mock_get_secret_str, mock_client, mock_cache
|
|
):
|
|
"""Test that WATSONX_ZENAPIKEY is used when it's the only key available."""
|
|
# Setup mocks
|
|
mock_cache.get_cache.return_value = None # Cache miss
|
|
mock_get_secret_str.side_effect = lambda key: (
|
|
"zen-api-key-12345" if key == "WATSONX_ZENAPIKEY" else None
|
|
)
|
|
|
|
mock_response = MagicMock()
|
|
mock_response.json.return_value = {
|
|
"access_token": "test-token-12345",
|
|
"expires_in": 3600,
|
|
}
|
|
mock_response.raise_for_status = MagicMock()
|
|
mock_client.post.return_value = mock_response
|
|
|
|
# Call function without api_key parameter
|
|
result = generate_iam_token()
|
|
|
|
# Verify get_secret_str was called with correct keys in order
|
|
# Note: get_watsonx_iam_url() also calls get_secret_str("WATSONX_IAM_URL")
|
|
calls = [
|
|
call[0][0]
|
|
for call in mock_get_secret_str.call_args_list
|
|
if call[0][0] != "WATSONX_IAM_URL"
|
|
]
|
|
assert "WX_API_KEY" in calls
|
|
assert "WATSONX_API_KEY" in calls
|
|
assert "WATSONX_APIKEY" in calls
|
|
assert "WATSONX_ZENAPIKEY" in calls
|
|
|
|
# Verify the token was generated using WATSONX_ZENAPIKEY
|
|
assert result == "test-token-12345"
|
|
mock_client.post.assert_called_once()
|
|
call_kwargs = mock_client.post.call_args
|
|
assert call_kwargs.kwargs["data"]["apikey"] == "zen-api-key-12345"
|
|
|
|
@patch("litellm.llms.watsonx.common_utils.iam_token_cache")
|
|
@patch("litellm.llms.watsonx.common_utils.litellm.module_level_client")
|
|
@patch("litellm.llms.watsonx.common_utils.get_secret_str")
|
|
def test_generate_iam_token_api_key_priority_order(
|
|
self, mock_get_secret_str, mock_client, mock_cache
|
|
):
|
|
"""Test that API keys are checked in the correct priority order."""
|
|
# Setup mocks
|
|
mock_cache.get_cache.return_value = None # Cache miss
|
|
|
|
# Test priority: WX_API_KEY > WATSONX_API_KEY > WATSONX_APIKEY > WATSONX_ZENAPIKEY
|
|
test_cases = [
|
|
# (env_keys_set, expected_key_used, expected_calls)
|
|
(
|
|
{"WX_API_KEY": "wx-key"},
|
|
"wx-key",
|
|
["WX_API_KEY"], # Should stop after first call
|
|
),
|
|
(
|
|
{"WATSONX_API_KEY": "watsonx-api-key"},
|
|
"watsonx-api-key",
|
|
["WX_API_KEY", "WATSONX_API_KEY"], # Should check WX_API_KEY first, then WATSONX_API_KEY
|
|
),
|
|
(
|
|
{"WATSONX_APIKEY": "watsonx-apikey"},
|
|
"watsonx-apikey",
|
|
["WX_API_KEY", "WATSONX_API_KEY", "WATSONX_APIKEY"],
|
|
),
|
|
(
|
|
{"WATSONX_ZENAPIKEY": "watsonx-zenapikey"},
|
|
"watsonx-zenapikey",
|
|
["WX_API_KEY", "WATSONX_API_KEY", "WATSONX_APIKEY", "WATSONX_ZENAPIKEY"],
|
|
),
|
|
# Test that higher priority keys take precedence
|
|
(
|
|
{
|
|
"WX_API_KEY": "wx-key",
|
|
"WATSONX_ZENAPIKEY": "zen-key",
|
|
},
|
|
"wx-key",
|
|
["WX_API_KEY"], # Should stop after first call
|
|
),
|
|
(
|
|
{
|
|
"WATSONX_API_KEY": "watsonx-api-key",
|
|
"WATSONX_ZENAPIKEY": "zen-key",
|
|
},
|
|
"watsonx-api-key",
|
|
["WX_API_KEY", "WATSONX_API_KEY"], # Should stop after WATSONX_API_KEY
|
|
),
|
|
(
|
|
{
|
|
"WATSONX_APIKEY": "watsonx-apikey",
|
|
"WATSONX_ZENAPIKEY": "zen-key",
|
|
},
|
|
"watsonx-apikey",
|
|
["WX_API_KEY", "WATSONX_API_KEY", "WATSONX_APIKEY"],
|
|
),
|
|
]
|
|
|
|
for env_keys, expected_key, expected_calls in test_cases:
|
|
mock_get_secret_str.reset_mock()
|
|
mock_client.reset_mock()
|
|
mock_cache.reset_mock()
|
|
|
|
# Configure mock to return values based on env_keys
|
|
def get_secret_side_effect(key):
|
|
return env_keys.get(key)
|
|
|
|
mock_get_secret_str.side_effect = get_secret_side_effect
|
|
|
|
mock_response = MagicMock()
|
|
mock_response.json.return_value = {
|
|
"access_token": "test-token",
|
|
"expires_in": 3600,
|
|
}
|
|
mock_response.raise_for_status = MagicMock()
|
|
mock_client.post.return_value = mock_response
|
|
|
|
# Call function
|
|
result = generate_iam_token()
|
|
|
|
# Verify the correct key was used
|
|
call_kwargs = mock_client.post.call_args
|
|
assert (
|
|
call_kwargs.kwargs["data"]["apikey"] == expected_key
|
|
), f"Expected {expected_key} but got {call_kwargs.kwargs['data']['apikey']} for env_keys: {env_keys}"
|
|
|
|
# Verify get_secret_str was called with expected keys (checking short-circuit behavior)
|
|
# Note: get_watsonx_iam_url() also calls get_secret_str("WATSONX_IAM_URL"), so we filter that out
|
|
actual_calls = [
|
|
call[0][0]
|
|
for call in mock_get_secret_str.call_args_list
|
|
if call[0][0] != "WATSONX_IAM_URL"
|
|
]
|
|
assert (
|
|
actual_calls == expected_calls
|
|
), f"Expected calls {expected_calls} but got {actual_calls} for env_keys: {env_keys}"
|
|
|
|
@patch("litellm.llms.watsonx.common_utils.iam_token_cache")
|
|
@patch("litellm.llms.watsonx.common_utils.litellm.module_level_client")
|
|
@patch("litellm.llms.watsonx.common_utils.get_secret_str")
|
|
def test_generate_iam_token_with_direct_api_key(
|
|
self, mock_get_secret_str, mock_client, mock_cache
|
|
):
|
|
"""Test that when api_key is passed directly, it's used instead of environment variables."""
|
|
# Setup mocks
|
|
mock_cache.get_cache.return_value = None # Cache miss
|
|
mock_get_secret_str.return_value = "env-key-should-not-be-used"
|
|
|
|
mock_response = MagicMock()
|
|
mock_response.json.return_value = {
|
|
"access_token": "test-token-12345",
|
|
"expires_in": 3600,
|
|
}
|
|
mock_response.raise_for_status = MagicMock()
|
|
mock_client.post.return_value = mock_response
|
|
|
|
# Call function with direct api_key
|
|
direct_key = "direct-api-key-12345"
|
|
result = generate_iam_token(api_key=direct_key)
|
|
|
|
# Verify get_secret_str was NOT called for API keys (since api_key was provided)
|
|
# Note: get_watsonx_iam_url() calls get_secret_str("WATSONX_IAM_URL"), which is expected
|
|
api_key_calls = [
|
|
call[0][0]
|
|
for call in mock_get_secret_str.call_args_list
|
|
if call[0][0] not in ["WATSONX_IAM_URL"]
|
|
]
|
|
assert (
|
|
len(api_key_calls) == 0
|
|
), f"Expected no API key calls but got {api_key_calls}"
|
|
|
|
# Verify the direct key was used
|
|
assert result == "test-token-12345"
|
|
call_kwargs = mock_client.post.call_args
|
|
assert call_kwargs.kwargs["data"]["apikey"] == direct_key
|
|
|
|
@patch("litellm.llms.watsonx.common_utils.iam_token_cache")
|
|
@patch("litellm.llms.watsonx.common_utils.get_secret_str")
|
|
def test_generate_iam_token_no_api_key_raises_error(
|
|
self, mock_get_secret_str, mock_cache
|
|
):
|
|
"""Test that ValueError is raised when no API key is available."""
|
|
# Setup mocks
|
|
mock_cache.get_cache.return_value = None # Cache miss
|
|
mock_get_secret_str.return_value = None # No keys available
|
|
|
|
# Call function without api_key and expect ValueError
|
|
with pytest.raises(ValueError, match="API key is required"):
|
|
generate_iam_token()
|
|
|
|
# Verify get_secret_str was called for all possible API keys
|
|
# Note: get_watsonx_iam_url() also calls get_secret_str("WATSONX_IAM_URL")
|
|
calls = [
|
|
call[0][0]
|
|
for call in mock_get_secret_str.call_args_list
|
|
if call[0][0] != "WATSONX_IAM_URL"
|
|
]
|
|
assert "WX_API_KEY" in calls
|
|
assert "WATSONX_API_KEY" in calls
|
|
assert "WATSONX_APIKEY" in calls
|
|
assert "WATSONX_ZENAPIKEY" in calls
|
|
|
|
@patch("litellm.llms.watsonx.common_utils.iam_token_cache")
|
|
@patch("litellm.llms.watsonx.common_utils.litellm.module_level_client")
|
|
@patch("litellm.llms.watsonx.common_utils.get_secret_str")
|
|
def test_generate_iam_token_uses_cache(
|
|
self, mock_get_secret_str, mock_client, mock_cache
|
|
):
|
|
"""Test that cached token is returned when available."""
|
|
# Setup mocks
|
|
cached_token = "cached-token-12345"
|
|
mock_cache.get_cache.return_value = cached_token
|
|
|
|
# Call function
|
|
result = generate_iam_token()
|
|
|
|
# Verify cached token was returned
|
|
assert result == cached_token
|
|
|
|
# Verify get_secret_str and client.post were NOT called (cache hit)
|
|
mock_get_secret_str.assert_not_called()
|
|
mock_client.post.assert_not_called()
|