scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-06 08:16:43 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 78

fix env var leakage

Browse source
This commit is contained in:
An Tang 2026-02-20 11:10:07 -08:00
parent be885ad236
commit 08a561f78a
1 changed files with 54 additions and 38 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

92
tests/test_litellm/llms/bedrock/test_base_aws_llm.py
View file

@ -555,9 +555,15 @@ def test_different_roles_without_session_names_should_not_share_cache():
)
def test_eks_irsa_ambient_credentials_used(role_kwargs, expected_client_kwargs):
"""
When only aws_role_name and aws_session_name are passed (no explicit creds),
boto3.client("sts", **expected) is called with no credential kwargs.
Test that in EKS/IRSA environments, ambient credentials are used when no explicit keys provided.
This allows web identity tokens to work automatically.
"""
# Isolate from ambient AWS_REGION/AWS_DEFAULT_REGION so no_region_or_endpoint is deterministic
env_without_aws_region = {
k: v
for k, v in os.environ.items()
if k not in ("AWS_REGION", "AWS_DEFAULT_REGION")
}
base_aws_llm = BaseAWSLLM()
mock_expiry = MagicMock()
mock_expiry.tzinfo = timezone.utc
@ -575,22 +581,25 @@ def test_eks_irsa_ambient_credentials_used(role_kwargs, expected_client_kwargs):
mock_sts_client = MagicMock()
mock_sts_client.assume_role.return_value = mock_sts_response
with patch("boto3.client", return_value=mock_sts_client) as mock_boto3_client:
credentials, ttl = base_aws_llm._auth_with_aws_role(
aws_access_key_id=None,
aws_secret_access_key=None,
aws_session_token=None,
aws_role_name="arn:aws:iam::2222222222222:role/LitellmEvalBedrockRole",
aws_session_name="test-session",
**role_kwargs,
)
mock_boto3_client.assert_called_once_with("sts", **expected_client_kwargs)
mock_sts_client.assume_role.assert_called_once_with(
RoleArn="arn:aws:iam::2222222222222:role/LitellmEvalBedrockRole",
RoleSessionName="test-session",
)
assert credentials.access_key == "assumed-access-key"
assert ttl is not None
with patch.dict(os.environ, env_without_aws_region, clear=True):
with patch("boto3.client", return_value=mock_sts_client) as mock_boto3_client:
credentials, ttl = base_aws_llm._auth_with_aws_role(
aws_access_key_id=None,
aws_secret_access_key=None,
aws_session_token=None,
aws_role_name="arn:aws:iam::2222222222222:role/LitellmEvalBedrockRole",
aws_session_name="test-session",
**role_kwargs,
)
mock_boto3_client.assert_called_once_with(
"sts", **expected_client_kwargs
)
mock_sts_client.assume_role.assert_called_once_with(
RoleArn="arn:aws:iam::2222222222222:role/LitellmEvalBedrockRole",
RoleSessionName="test-session",
)
assert credentials.access_key == "assumed-access-key"
assert ttl is not None
@pytest.mark.parametrize(
@ -632,9 +641,13 @@ def test_explicit_credentials_used_when_provided(role_kwargs, expected_client_kw
"""
Test that explicit credentials are used when provided (non-EKS/IRSA scenario).
"""
# Isolate from ambient AWS_REGION/AWS_DEFAULT_REGION so no_region_or_endpoint is deterministic
env_without_aws_region = {
k: v
for k, v in os.environ.items()
if k not in ("AWS_REGION", "AWS_DEFAULT_REGION")
}
base_aws_llm = BaseAWSLLM()
# Mock the STS response with proper expiration handling
mock_expiry = MagicMock()
mock_expiry.tzinfo = timezone.utc
# Create a timedelta object that returns 3600 when total_seconds() is called
@ -652,24 +665,27 @@ def test_explicit_credentials_used_when_provided(role_kwargs, expected_client_kw
mock_sts_client = MagicMock()
mock_sts_client.assume_role.return_value = mock_sts_response
with patch("boto3.client", return_value=mock_sts_client) as mock_boto3_client:
credentials, ttl = base_aws_llm._auth_with_aws_role(
aws_access_key_id="explicit-access-key",
aws_secret_access_key="explicit-secret-key",
aws_session_token="assumed-session-token",
aws_role_name="arn:aws:iam::2222222222222:role/LitellmEvalBedrockRole",
aws_session_name="test-session",
**role_kwargs
)
mock_boto3_client.assert_called_once_with("sts", **expected_client_kwargs)
mock_sts_client.assume_role.assert_called_once_with(
RoleArn="arn:aws:iam::2222222222222:role/LitellmEvalBedrockRole",
RoleSessionName="test-session"
)
assert credentials.access_key == "assumed-access-key"
assert credentials.secret_key == "assumed-secret-key"
assert credentials.token == "assumed-session-token"
assert ttl is not None
with patch.dict(os.environ, env_without_aws_region, clear=True):
with patch("boto3.client", return_value=mock_sts_client) as mock_boto3_client:
credentials, ttl = base_aws_llm._auth_with_aws_role(
aws_access_key_id="explicit-access-key",
aws_secret_access_key="explicit-secret-key",
aws_session_token="assumed-session-token",
aws_role_name="arn:aws:iam::2222222222222:role/LitellmEvalBedrockRole",
aws_session_name="test-session",
**role_kwargs,
)
mock_boto3_client.assert_called_once_with(
"sts", **expected_client_kwargs
)
mock_sts_client.assume_role.assert_called_once_with(
RoleArn="arn:aws:iam::2222222222222:role/LitellmEvalBedrockRole",
RoleSessionName="test-session",
)
assert credentials.access_key == "assumed-access-key"
assert credentials.secret_key == "assumed-secret-key"
assert credentials.token == "assumed-session-token"
assert ttl is not None
def test_partial_credentials_still_use_ambient():