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 178

update tests to include case with aws_region_name for _auth_with_aws_role

Browse source
This commit is contained in:
An Tang 2026-02-19 18:06:35 -08:00
parent c51696ddda
commit 5584f2d612
1 changed files with 24 additions and 6 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -548,9 +548,6 @@ def test_eks_irsa_ambient_credentials_used():
"""
base_aws_llm = BaseAWSLLM()
# Mock the boto3 STS client
mock_sts_client = MagicMock()
# Mock the STS response with proper expiration handling
mock_expiry = MagicMock()
mock_expiry.tzinfo = timezone.utc
@ -568,6 +565,9 @@ def test_eks_irsa_ambient_credentials_used():
"Expiration": mock_expiry,
}
}
# Case 1: only aws_role_name and aws_session_name (no aws_region_name)
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:
@ -588,13 +588,31 @@ def test_eks_irsa_ambient_credentials_used():
# Should call assume_role
mock_sts_client.assume_role.assert_called_once_with(
RoleArn="arn:aws:iam::2222222222222:role/LitellmEvalBedrockRole",
RoleSessionName="test-session"
RoleSessionName="test-session",
)
# Verify credentials are returned correctly
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
# Case 2: aws_role_name, aws_session_name, and aws_region_name (regional STS)
mock_sts_client2 = MagicMock()
mock_sts_client2.assume_role.return_value = mock_sts_response
with patch("boto3.client", return_value=mock_sts_client2) 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",
aws_region_name="us-east-1",
)
mock_boto3_client.assert_called_once_with("sts", region_name="us-east-1", verify=True)
mock_sts_client2.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