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 a38e0c0b0ac..18fc7c6173e 100644 --- a/tests/test_litellm/llms/bedrock/test_base_aws_llm.py +++ b/tests/test_litellm/llms/bedrock/test_base_aws_llm.py @@ -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():