diff --git a/tests/test_litellm/llms/bedrock/test_base_aws_llm.py b/tests/test_litellm/llms/bedrock/test_base_aws_llm.py index cf9fee6bacf..335f057edbb 100644 --- a/tests/test_litellm/llms/bedrock/test_base_aws_llm.py +++ b/tests/test_litellm/llms/bedrock/test_base_aws_llm.py @@ -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