ai-workspace-services/litellm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge pull request #13753 from RichardoC/fix-bedrock-irsa-role-assumption

Browse Source 
fix: role chaining and session name with webauthentication for aws bedrock
This commit is contained in:
Krish Dholakia 2025-08-22 07:20:27 -07:00 committed by GitHub
commit dcad39c00f
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 728 additions and 16 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

179
litellm/llms/bedrock/base_aws_llm.py
View File

@ -179,15 +179,32 @@ class BaseAWSLLM:
aws_sts_endpoint=aws_sts_endpoint,
)
elif aws_role_name is not None:
# If aws_session_name is not provided, generate a default one
if aws_session_name is None:
aws_session_name = f"litellm-session-{int(datetime.now().timestamp())}"
credentials, _cache_ttl = self._auth_with_aws_role(
aws_access_key_id=aws_access_key_id,
aws_secret_access_key=aws_secret_access_key,
aws_role_name=aws_role_name,
aws_session_name=aws_session_name,
)
# Check if we're in IRSA and trying to assume the same role we already have
current_role_arn = os.getenv("AWS_ROLE_ARN")
web_identity_token_file = os.getenv("AWS_WEB_IDENTITY_TOKEN_FILE")
# In IRSA environments, we should skip role assumption if we're already running as the target role
# This is true when:
# 1. We have AWS_ROLE_ARN set (current role)
# 2. We have AWS_WEB_IDENTITY_TOKEN_FILE set (IRSA environment)
# 3. The current role matches the requested role
if (current_role_arn and web_identity_token_file and
current_role_arn == aws_role_name):
verbose_logger.debug("Using IRSA same-role optimization: calling _auth_with_env_vars")
# We're already running as this role via IRSA, no need to assume it again
# Use the default boto3 credentials (which will use the IRSA credentials)
credentials, _cache_ttl = self._auth_with_env_vars()
else:
verbose_logger.debug("Using role assumption: calling _auth_with_aws_role")
# If aws_session_name is not provided, generate a default one
if aws_session_name is None:
aws_session_name = f"litellm-session-{int(datetime.now().timestamp())}"
credentials, _cache_ttl = self._auth_with_aws_role(
aws_access_key_id=aws_access_key_id,
aws_secret_access_key=aws_secret_access_key,
aws_role_name=aws_role_name,
aws_session_name=aws_session_name,
)
elif aws_profile_name is not None: ### CHECK SESSION ###
credentials, _cache_ttl = self._auth_with_aws_profile(aws_profile_name)
@ -446,6 +463,92 @@ class BaseAWSLLM:
iam_creds = session.get_credentials()
return iam_creds, self._get_default_ttl_for_boto3_credentials()
def _handle_irsa_cross_account(self, irsa_role_arn: str, aws_role_name: str,
aws_session_name: str, region: str, web_identity_token_file: str) -> dict:
"""Handle cross-account role assumption for IRSA."""
import boto3
verbose_logger.debug("Cross-account role assumption detected")
# Read the web identity token
with open(web_identity_token_file, 'r') as f:
web_identity_token = f.read().strip()
# Create an STS client without credentials
with tracer.trace("boto3.client(sts) for manual IRSA"):
sts_client = boto3.client('sts', region_name=region)
# Manually assume the IRSA role with the session name
verbose_logger.debug(f"Manually assuming IRSA role {irsa_role_arn} with session {aws_session_name}")
irsa_response = sts_client.assume_role_with_web_identity(
RoleArn=irsa_role_arn,
RoleSessionName=aws_session_name,
WebIdentityToken=web_identity_token
)
# Extract the credentials from the IRSA assumption
irsa_creds = irsa_response["Credentials"]
# Create a new STS client with the IRSA credentials
with tracer.trace("boto3.client(sts) with manual IRSA credentials"):
sts_client_with_creds = boto3.client(
'sts',
region_name=region,
aws_access_key_id=irsa_creds["AccessKeyId"],
aws_secret_access_key=irsa_creds["SecretAccessKey"],
aws_session_token=irsa_creds["SessionToken"]
)
# Get current caller identity for debugging
try:
caller_identity = sts_client_with_creds.get_caller_identity()
verbose_logger.debug(f"Current identity after manual IRSA assumption: {caller_identity.get('Arn', 'unknown')}")
except Exception as e:
verbose_logger.debug(f"Failed to get caller identity: {e}")
# Now assume the target role
verbose_logger.debug(f"Attempting to assume target role: {aws_role_name} with session: {aws_session_name}")
return sts_client_with_creds.assume_role(
RoleArn=aws_role_name, RoleSessionName=aws_session_name
)
def _handle_irsa_same_account(self, aws_role_name: str, aws_session_name: str, region: str) -> dict:
"""Handle same-account role assumption for IRSA."""
import boto3
verbose_logger.debug("Same account role assumption, using automatic IRSA")
with tracer.trace("boto3.client(sts) with automatic IRSA"):
sts_client = boto3.client("sts", region_name=region)
# Get current caller identity for debugging
try:
caller_identity = sts_client.get_caller_identity()
verbose_logger.debug(f"Current IRSA identity: {caller_identity.get('Arn', 'unknown')}")
except Exception as e:
verbose_logger.debug(f"Failed to get caller identity: {e}")
# Assume the role
verbose_logger.debug(f"Attempting to assume role: {aws_role_name} with session: {aws_session_name}")
return sts_client.assume_role(
RoleArn=aws_role_name, RoleSessionName=aws_session_name
)
def _extract_credentials_and_ttl(self, sts_response: dict) -> Tuple[Credentials, Optional[int]]:
"""Extract credentials and TTL from STS response."""
from botocore.credentials import Credentials
sts_credentials = sts_response["Credentials"]
credentials = Credentials(
access_key=sts_credentials["AccessKeyId"],
secret_key=sts_credentials["SecretAccessKey"],
token=sts_credentials["SessionToken"],
)
expiration_time = sts_credentials["Expiration"]
ttl = int((expiration_time - datetime.now(expiration_time.tzinfo)).total_seconds())
return credentials, ttl
@tracer.wrap()
def _auth_with_aws_role(
self,
@ -460,12 +563,58 @@ class BaseAWSLLM:
import boto3
from botocore.credentials import Credentials
with tracer.trace("boto3.client(sts)"):
sts_client = boto3.client(
"sts",
aws_access_key_id=aws_access_key_id, # [OPTIONAL]
aws_secret_access_key=aws_secret_access_key, # [OPTIONAL]
)
# Check if we're in an EKS/IRSA environment
web_identity_token_file = os.getenv("AWS_WEB_IDENTITY_TOKEN_FILE")
irsa_role_arn = os.getenv("AWS_ROLE_ARN")
# If we have IRSA environment variables and no explicit credentials,
# we need to use the web identity token flow
if (web_identity_token_file and irsa_role_arn and
aws_access_key_id is None and aws_secret_access_key is None):
# For cross-account role assumption with specific session names,
# we need to manually assume the IRSA role first with the correct session name
verbose_logger.debug(f"IRSA detected: using web identity token from {web_identity_token_file}")
try:
# Get region from environment
region = os.getenv("AWS_REGION") or os.getenv("AWS_DEFAULT_REGION") or "us-east-1"
# Check if we need to do cross-account role assumption
if aws_role_name != irsa_role_arn:
sts_response = self._handle_irsa_cross_account(
irsa_role_arn, aws_role_name, aws_session_name, region, web_identity_token_file
)
else:
sts_response = self._handle_irsa_same_account(
aws_role_name, aws_session_name, region
)
return self._extract_credentials_and_ttl(sts_response)
except Exception as e:
verbose_logger.debug(f"Failed to assume role via IRSA: {e}")
if "AccessDenied" in str(e) and "is not authorized to perform: sts:AssumeRole" in str(e):
# Provide a more helpful error message for trust policy issues
verbose_logger.error(
f"Access denied when trying to assume role {aws_role_name}. "
f"Please ensure the trust policy of {aws_role_name} allows "
f"the current role to assume it. Current identity: check logs with verbose mode."
)
# Re-raise the exception instead of falling through
raise
# In EKS/IRSA environments, use ambient credentials (no explicit keys needed)
# This allows the web identity token to work automatically
if aws_access_key_id is None and aws_secret_access_key is None:
with tracer.trace("boto3.client(sts)"):
sts_client = boto3.client("sts")
else:
with tracer.trace("boto3.client(sts)"):
sts_client = boto3.client(
"sts",
aws_access_key_id=aws_access_key_id,
aws_secret_access_key=aws_secret_access_key,
)
sts_response = sts_client.assume_role(
RoleArn=aws_role_name, RoleSessionName=aws_session_name

565
tests/test_litellm/llms/bedrock/test_base_aws_llm.py
View File

@ -10,7 +10,7 @@ sys.path.insert(
) # Adds the parent directory to the system path
from datetime import datetime, timezone
from datetime import datetime, timedelta, timezone
from typing import Any, Dict
from unittest.mock import MagicMock, patch
@ -479,3 +479,566 @@ def test_role_assumption_without_session_name():
# Should only be called once due to caching
assert mock_sts_client.assume_role.call_count == 1
def test_cache_keys_are_different_for_different_roles():
"""
Test that cache keys are different for different AWS roles.
This ensures that credentials for different roles don't get mixed up.
"""
base_aws_llm = BaseAWSLLM()
# Create arguments for two different roles
args1 = {
"aws_access_key_id": None,
"aws_secret_access_key": None,
"aws_role_name": "arn:aws:iam::1111111111111:role/LitellmRole",
"aws_session_name": "test-session-1"
}
args2 = {
"aws_access_key_id": None,
"aws_secret_access_key": None,
"aws_role_name": "arn:aws:iam::2222222222222:role/LitellmEvalBedrockRole",
"aws_session_name": "test-session-2"
}
# Generate cache keys
cache_key1 = base_aws_llm.get_cache_key(args1)
cache_key2 = base_aws_llm.get_cache_key(args2)
# Cache keys should be different because the role names are different
assert cache_key1 != cache_key2
def test_different_roles_without_session_names_should_not_share_cache():
"""
Test that different roles with auto-generated session names don't share cache.
This was the original issue where cache keys were the same for different roles.
"""
base_aws_llm = BaseAWSLLM()
# Create arguments for two different roles without session names
args1 = {
"aws_access_key_id": None,
"aws_secret_access_key": None,
"aws_role_name": "arn:aws:iam::1111111111111:role/LitellmRole",
"aws_session_name": None
}
args2 = {
"aws_access_key_id": None,
"aws_secret_access_key": None,
"aws_role_name": "arn:aws:iam::2222222222222:role/LitellmEvalBedrockRole",
"aws_session_name": None
}
# Generate cache keys
cache_key1 = base_aws_llm.get_cache_key(args1)
cache_key2 = base_aws_llm.get_cache_key(args2)
# Cache keys should be different because the role names are different
assert cache_key1 != cache_key2
def test_eks_irsa_ambient_credentials_used():
"""
Test that in EKS/IRSA environments, ambient credentials are used when no explicit keys provided.
This allows web identity tokens to work automatically.
"""
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
current_time = datetime.now(timezone.utc)
# Create a timedelta object that returns 3600 when total_seconds() is called
time_diff = MagicMock()
time_diff.total_seconds.return_value = 3600
mock_expiry.__sub__ = MagicMock(return_value=time_diff)
mock_sts_response = {
"Credentials": {
"AccessKeyId": "assumed-access-key",
"SecretAccessKey": "assumed-secret-key",
"SessionToken": "assumed-session-token",
"Expiration": mock_expiry,
}
}
mock_sts_client.assume_role.return_value = mock_sts_response
with patch("boto3.client", return_value=mock_sts_client) as mock_boto3_client:
# Call with no explicit credentials (EKS/IRSA scenario)
credentials, ttl = base_aws_llm._auth_with_aws_role(
aws_access_key_id=None,
aws_secret_access_key=None,
aws_role_name="arn:aws:iam::2222222222222:role/LitellmEvalBedrockRole",
aws_session_name="test-session"
)
# Should create STS client without explicit credentials (using ambient credentials)
mock_boto3_client.assert_called_once_with("sts")
# Should call assume_role
mock_sts_client.assume_role.assert_called_once_with(
RoleArn="arn:aws:iam::2222222222222:role/LitellmEvalBedrockRole",
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
def test_explicit_credentials_used_when_provided():
"""
Test that explicit credentials are used when provided (non-EKS/IRSA scenario).
"""
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
current_time = datetime.now(timezone.utc)
# Create a timedelta object that returns 3600 when total_seconds() is called
time_diff = MagicMock()
time_diff.total_seconds.return_value = 3600
mock_expiry.__sub__ = MagicMock(return_value=time_diff)
mock_sts_response = {
"Credentials": {
"AccessKeyId": "assumed-access-key",
"SecretAccessKey": "assumed-secret-key",
"SessionToken": "assumed-session-token",
"Expiration": mock_expiry,
}
}
mock_sts_client.assume_role.return_value = mock_sts_response
with patch("boto3.client", return_value=mock_sts_client) as mock_boto3_client:
# Call with explicit credentials
credentials, ttl = base_aws_llm._auth_with_aws_role(
aws_access_key_id="explicit-access-key",
aws_secret_access_key="explicit-secret-key",
aws_role_name="arn:aws:iam::2222222222222:role/LitellmEvalBedrockRole",
aws_session_name="test-session"
)
# Should create STS client with explicit credentials
mock_boto3_client.assert_called_once_with(
"sts",
aws_access_key_id="explicit-access-key",
aws_secret_access_key="explicit-secret-key",
)
# Should call assume_role
mock_sts_client.assume_role.assert_called_once_with(
RoleArn="arn:aws:iam::2222222222222:role/LitellmEvalBedrockRole",
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
def test_partial_credentials_still_use_ambient():
"""
Test that if only one credential is provided, we still use ambient credentials.
This handles edge cases where configuration might be incomplete.
"""
base_aws_llm = BaseAWSLLM()
# Mock the boto3 STS client
mock_sts_client = MagicMock()
# Mock the STS response
mock_expiry = MagicMock()
mock_expiry.tzinfo = timezone.utc
time_diff = MagicMock()
time_diff.total_seconds.return_value = 3600
mock_expiry.__sub__ = MagicMock(return_value=time_diff)
mock_sts_response = {
"Credentials": {
"AccessKeyId": "assumed-access-key",
"SecretAccessKey": "assumed-secret-key",
"SessionToken": "assumed-session-token",
"Expiration": mock_expiry,
}
}
mock_sts_client.assume_role.return_value = mock_sts_response
with patch("boto3.client", return_value=mock_sts_client) as mock_boto3_client:
# Call with only access key (missing secret key)
credentials, ttl = base_aws_llm._auth_with_aws_role(
aws_access_key_id="AKIAEXAMPLE",
aws_secret_access_key=None,
aws_role_name="arn:aws:iam::2222222222222:role/LitellmEvalBedrockRole",
aws_session_name="test-session"
)
# Should still pass partial credentials to boto3.client
mock_boto3_client.assert_called_once_with(
"sts",
aws_access_key_id="AKIAEXAMPLE",
aws_secret_access_key=None
)
# Should still call assume_role
mock_sts_client.assume_role.assert_called_once_with(
RoleArn="arn:aws:iam::2222222222222:role/LitellmEvalBedrockRole",
RoleSessionName="test-session"
)
def test_cross_account_role_assumption():
"""
Test assuming a role in a different AWS account (common in multi-account setups).
"""
base_aws_llm = BaseAWSLLM()
# Mock the boto3 STS client
mock_sts_client = MagicMock()
# Mock the STS response for cross-account role
mock_expiry = MagicMock()
mock_expiry.tzinfo = timezone.utc
time_diff = MagicMock()
time_diff.total_seconds.return_value = 3600
mock_expiry.__sub__ = MagicMock(return_value=time_diff)
mock_sts_response = {
"Credentials": {
"AccessKeyId": "cross-account-access-key",
"SecretAccessKey": "cross-account-secret-key",
"SessionToken": "cross-account-session-token",
"Expiration": mock_expiry,
}
}
mock_sts_client.assume_role.return_value = mock_sts_response
with patch("boto3.client", return_value=mock_sts_client) as mock_boto3_client:
# Assume role in different account (EKS/IRSA scenario)
credentials, ttl = base_aws_llm._auth_with_aws_role(
aws_access_key_id=None,
aws_secret_access_key=None,
aws_role_name="arn:aws:iam::999999999999:role/CrossAccountRole",
aws_session_name="cross-account-session"
)
# Should use ambient credentials
mock_boto3_client.assert_called_once_with("sts")
# Should call assume_role with cross-account role
mock_sts_client.assume_role.assert_called_once_with(
RoleArn="arn:aws:iam::999999999999:role/CrossAccountRole",
RoleSessionName="cross-account-session"
)
# Verify cross-account credentials are returned
assert credentials.access_key == "cross-account-access-key"
assert credentials.secret_key == "cross-account-secret-key"
assert credentials.token == "cross-account-session-token"
assert ttl is not None
def test_role_assumption_with_custom_session_name():
"""
Test role assumption with a custom session name.
"""
base_aws_llm = BaseAWSLLM()
# Mock the boto3 STS client
mock_sts_client = MagicMock()
# Mock the STS response
mock_expiry = MagicMock()
mock_expiry.tzinfo = timezone.utc
time_diff = MagicMock()
time_diff.total_seconds.return_value = 3600
mock_expiry.__sub__ = MagicMock(return_value=time_diff)
mock_sts_response = {
"Credentials": {
"AccessKeyId": "custom-session-access-key",
"SecretAccessKey": "custom-session-secret-key",
"SessionToken": "custom-session-token",
"Expiration": mock_expiry,
}
}
mock_sts_client.assume_role.return_value = mock_sts_response
with patch("boto3.client", return_value=mock_sts_client):
# Use custom session name
credentials, ttl = base_aws_llm._auth_with_aws_role(
aws_access_key_id=None,
aws_secret_access_key=None,
aws_role_name="arn:aws:iam::1111111111111:role/LitellmRole",
aws_session_name="evals-bedrock-session"
)
# Should call assume_role with custom session name
mock_sts_client.assume_role.assert_called_once_with(
RoleArn="arn:aws:iam::1111111111111:role/LitellmRole",
RoleSessionName="evals-bedrock-session"
)
# Verify credentials are returned
assert credentials.access_key == "custom-session-access-key"
assert credentials.secret_key == "custom-session-secret-key"
assert credentials.token == "custom-session-token"
def test_role_assumption_ttl_calculation():
"""
Test that TTL is calculated correctly from STS response expiration.
"""
base_aws_llm = BaseAWSLLM()
# Mock the boto3 STS client
mock_sts_client = MagicMock()
# Create a real datetime for expiration (1 hour from now)
expiration_time = datetime.now(timezone.utc) + timedelta(hours=1)
mock_sts_response = {
"Credentials": {
"AccessKeyId": "ttl-test-access-key",
"SecretAccessKey": "ttl-test-secret-key",
"SessionToken": "ttl-test-session-token",
"Expiration": expiration_time,
}
}
mock_sts_client.assume_role.return_value = mock_sts_response
with patch("boto3.client", return_value=mock_sts_client):
credentials, ttl = base_aws_llm._auth_with_aws_role(
aws_access_key_id=None,
aws_secret_access_key=None,
aws_role_name="arn:aws:iam::1111111111111:role/LitellmRole",
aws_session_name="ttl-test-session"
)
# TTL should be approximately 3540 seconds (1 hour - 60 second buffer)
assert ttl is not None
assert 3500 <= ttl <= 3600 # Allow some variance for test execution time
def test_role_assumption_error_handling():
"""
Test that role assumption errors are properly propagated.
"""
base_aws_llm = BaseAWSLLM()
# Mock the boto3 STS client to raise an exception
mock_sts_client = MagicMock()
mock_sts_client.assume_role.side_effect = Exception("AccessDenied: User is not authorized to perform sts:AssumeRole")
with patch("boto3.client", return_value=mock_sts_client):
# Should raise the exception
with pytest.raises(Exception) as exc_info:
base_aws_llm._auth_with_aws_role(
aws_access_key_id=None,
aws_secret_access_key=None,
aws_role_name="arn:aws:iam::1111111111111:role/UnauthorizedRole",
aws_session_name="error-test-session"
)
assert "AccessDenied" in str(exc_info.value)
def test_multiple_role_assumptions_in_sequence():
"""
Test that multiple role assumptions work correctly in sequence.
This simulates the scenario where different models use different roles.
"""
base_aws_llm = BaseAWSLLM()
# Mock the boto3 STS client
mock_sts_client = MagicMock()
# Mock different responses for different roles
mock_expiry = MagicMock()
mock_expiry.tzinfo = timezone.utc
time_diff = MagicMock()
time_diff.total_seconds.return_value = 3600
mock_expiry.__sub__ = MagicMock(return_value=time_diff)
# First role response
mock_sts_response1 = {
"Credentials": {
"AccessKeyId": "role1-access-key",
"SecretAccessKey": "role1-secret-key",
"SessionToken": "role1-session-token",
"Expiration": mock_expiry,
}
}
# Second role response
mock_sts_response2 = {
"Credentials": {
"AccessKeyId": "role2-access-key",
"SecretAccessKey": "role2-secret-key",
"SessionToken": "role2-session-token",
"Expiration": mock_expiry,
}
}
# Configure mock to return different responses
mock_sts_client.assume_role.side_effect = [mock_sts_response1, mock_sts_response2]
with patch("boto3.client", return_value=mock_sts_client):
# First role assumption
credentials1, ttl1 = base_aws_llm._auth_with_aws_role(
aws_access_key_id=None,
aws_secret_access_key=None,
aws_role_name="arn:aws:iam::1111111111111:role/LitellmRole",
aws_session_name="session-1"
)
# Second role assumption
credentials2, ttl2 = base_aws_llm._auth_with_aws_role(
aws_access_key_id=None,
aws_secret_access_key=None,
aws_role_name="arn:aws:iam::2222222222222:role/LitellmEvalBedrockRole",
aws_session_name="session-2"
)
# Verify both role assumptions were made
assert mock_sts_client.assume_role.call_count == 2
# Verify first role credentials
assert credentials1.access_key == "role1-access-key"
assert credentials1.secret_key == "role1-secret-key"
assert credentials1.token == "role1-session-token"
# Verify second role credentials
assert credentials2.access_key == "role2-access-key"
assert credentials2.secret_key == "role2-secret-key"
assert credentials2.token == "role2-session-token"
def test_auth_with_aws_role_irsa_environment():
"""Test that _auth_with_aws_role detects and uses IRSA environment variables"""
base_llm = BaseAWSLLM()
# Create a temporary file to simulate the web identity token
import tempfile
with tempfile.NamedTemporaryFile(mode='w', delete=False) as f:
f.write('test-web-identity-token')
token_file = f.name
try:
# Set IRSA environment variables
with patch.dict(os.environ, {
'AWS_WEB_IDENTITY_TOKEN_FILE': token_file,
'AWS_ROLE_ARN': 'arn:aws:iam::111111111111:role/eks-service-account-role',
'AWS_REGION': 'us-east-1'
}):
# Mock the boto3 STS client
mock_sts_client = MagicMock()
mock_assume_web_identity_response = {
'Credentials': {
'AccessKeyId': 'irsa-temp-access-key',
'SecretAccessKey': 'irsa-temp-secret-key',
'SessionToken': 'irsa-temp-session-token',
'Expiration': datetime.now() + timedelta(hours=1)
}
}
mock_assume_role_response = {
'Credentials': {
'AccessKeyId': 'irsa-access-key',
'SecretAccessKey': 'irsa-secret-key',
'SessionToken': 'irsa-session-token',
'Expiration': datetime.now() + timedelta(hours=1)
}
}
mock_sts_client.assume_role_with_web_identity.return_value = mock_assume_web_identity_response
mock_sts_client.assume_role.return_value = mock_assume_role_response
with patch('boto3.client', return_value=mock_sts_client) as mock_boto3_client:
# Call _auth_with_aws_role without explicit credentials
creds, ttl = base_llm._auth_with_aws_role(
aws_access_key_id=None,
aws_secret_access_key=None,
aws_role_name='arn:aws:iam::222222222222:role/target-role',
aws_session_name='test-session'
)
# Verify boto3.client was called multiple times
# First for manual IRSA, then with IRSA credentials
assert mock_boto3_client.call_count >= 2
# Verify assume_role_with_web_identity was called
mock_sts_client.assume_role_with_web_identity.assert_called_once_with(
RoleArn='arn:aws:iam::111111111111:role/eks-service-account-role',
RoleSessionName='test-session',
WebIdentityToken='test-web-identity-token'
)
# Verify assume_role was called with correct parameters
mock_sts_client.assume_role.assert_called_once_with(
RoleArn='arn:aws:iam::222222222222:role/target-role',
RoleSessionName='test-session'
)
# Verify the returned credentials
assert creds.access_key == 'irsa-access-key'
assert creds.secret_key == 'irsa-secret-key'
assert creds.token == 'irsa-session-token'
assert ttl > 0 # TTL should be positive
finally:
# Clean up the temporary file
os.unlink(token_file)
def test_auth_with_aws_role_same_role_irsa():
"""Test that when IRSA role matches the requested role, we skip assumption"""
base_llm = BaseAWSLLM()
# Set IRSA environment variables
with patch.dict(os.environ, {
'AWS_ROLE_ARN': 'arn:aws:iam::111111111111:role/LitellmRole',
'AWS_WEB_IDENTITY_TOKEN_FILE': '/var/run/secrets/eks.amazonaws.com/serviceaccount/token'
}):
# Mock the _auth_with_env_vars method
mock_creds = MagicMock()
mock_creds.access_key = 'irsa-access-key'
mock_creds.secret_key = 'irsa-secret-key'
mock_creds.token = 'irsa-session-token'
with patch.object(base_llm, '_auth_with_env_vars', return_value=(mock_creds, None)) as mock_env_auth:
# Call get_credentials instead of _auth_with_aws_role directly
# This tests the full flow
creds = base_llm.get_credentials(
aws_access_key_id=None,
aws_secret_access_key=None,
aws_role_name='arn:aws:iam::111111111111:role/LitellmRole', # Same as AWS_ROLE_ARN
aws_session_name='test-session',
aws_region_name='us-east-1'
)
# Verify it used the env vars auth (no role assumption)
mock_env_auth.assert_called_once()
# Verify the returned credentials
assert creds.access_key == 'irsa-access-key'