mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-01-10 00:28:31 +00:00
f36d48d57e
- Add clearer authentication method comments in documentation examples - Include security best practices for API key usage - Add CLI validation hints for AWS Bedrock API key format - Implement comprehensive error handling tests (timeout, rate limit, invalid key) - Add connection error and retry scenario test cases - Maintain consistency between English and Portuguese documentation Addresses AI code review feedback on PR #3126 Co-Authored-By: Jo\u00E3o <joao@crewai.com>
178 lines
7.1 KiB
Python
178 lines
7.1 KiB
Python
import os
|
|
import pytest
|
|
from unittest.mock import patch, MagicMock
|
|
from crewai import LLM
|
|
|
|
|
|
class TestBedrockAuthentication:
|
|
"""Test AWS Bedrock authentication methods."""
|
|
|
|
@patch.dict(os.environ, {
|
|
'AWS_ACCESS_KEY_ID': 'test-key-id',
|
|
'AWS_SECRET_ACCESS_KEY': 'test-secret-key',
|
|
'AWS_DEFAULT_REGION': 'us-east-1'
|
|
})
|
|
@patch('litellm.completion')
|
|
def test_bedrock_iam_authentication(self, mock_completion):
|
|
"""Test Bedrock with IAM role authentication."""
|
|
mock_completion.return_value = MagicMock()
|
|
mock_completion.return_value.choices = [MagicMock()]
|
|
mock_completion.return_value.choices[0].message.content = "Test response"
|
|
|
|
llm = LLM(model="bedrock/anthropic.claude-3-sonnet-20240229-v1:0")
|
|
result = llm.call("test message")
|
|
|
|
mock_completion.assert_called_once()
|
|
assert result == "Test response"
|
|
|
|
@patch.dict(os.environ, {
|
|
'AWS_BEARER_TOKEN_BEDROCK': 'test-api-key',
|
|
'AWS_DEFAULT_REGION': 'us-east-1'
|
|
})
|
|
@patch('litellm.completion')
|
|
def test_bedrock_api_key_authentication(self, mock_completion):
|
|
"""Test Bedrock with API key authentication."""
|
|
mock_completion.return_value = MagicMock()
|
|
mock_completion.return_value.choices = [MagicMock()]
|
|
mock_completion.return_value.choices[0].message.content = "Test response"
|
|
|
|
llm = LLM(model="bedrock/anthropic.claude-3-sonnet-20240229-v1:0")
|
|
result = llm.call("test message")
|
|
|
|
mock_completion.assert_called_once()
|
|
assert result == "Test response"
|
|
|
|
def test_bedrock_missing_credentials(self):
|
|
"""Test Bedrock fails gracefully with missing credentials."""
|
|
with patch.dict(os.environ, {}, clear=True):
|
|
llm = LLM(model="bedrock/anthropic.claude-3-sonnet-20240229-v1:0")
|
|
assert llm.model == "bedrock/anthropic.claude-3-sonnet-20240229-v1:0"
|
|
|
|
@patch.dict(os.environ, {
|
|
'AWS_BEARER_TOKEN_BEDROCK': 'test-api-key',
|
|
'AWS_DEFAULT_REGION': 'us-east-1'
|
|
})
|
|
@patch('litellm.completion')
|
|
def test_bedrock_api_key_with_streaming(self, mock_completion):
|
|
"""Test Bedrock API key authentication with streaming."""
|
|
mock_completion.return_value = iter([
|
|
MagicMock(choices=[MagicMock(delta=MagicMock(content="Test"))]),
|
|
MagicMock(choices=[MagicMock(delta=MagicMock(content=" response"))])
|
|
])
|
|
|
|
llm = LLM(model="bedrock/anthropic.claude-3-sonnet-20240229-v1:0")
|
|
result = list(llm.stream("test message"))
|
|
|
|
mock_completion.assert_called_once()
|
|
assert len(result) == 2
|
|
|
|
@patch.dict(os.environ, {
|
|
'AWS_ACCESS_KEY_ID': 'test-key-id',
|
|
'AWS_SECRET_ACCESS_KEY': 'test-secret-key',
|
|
'AWS_DEFAULT_REGION': 'us-east-1'
|
|
})
|
|
@patch('litellm.completion')
|
|
def test_bedrock_iam_with_custom_parameters(self, mock_completion):
|
|
"""Test Bedrock IAM authentication with custom parameters."""
|
|
mock_completion.return_value = MagicMock()
|
|
mock_completion.return_value.choices = [MagicMock()]
|
|
mock_completion.return_value.choices[0].message.content = "Test response"
|
|
|
|
llm = LLM(
|
|
model="bedrock/anthropic.claude-3-sonnet-20240229-v1:0",
|
|
temperature=0.7,
|
|
max_tokens=100
|
|
)
|
|
result = llm.call("test message")
|
|
|
|
mock_completion.assert_called_once()
|
|
call_args = mock_completion.call_args
|
|
assert call_args[1]['temperature'] == 0.7
|
|
assert call_args[1]['max_tokens'] == 100
|
|
assert result == "Test response"
|
|
|
|
@patch.dict(os.environ, {
|
|
'AWS_BEARER_TOKEN_BEDROCK': 'test-api-key',
|
|
'AWS_DEFAULT_REGION': 'us-west-2'
|
|
})
|
|
@patch('litellm.completion')
|
|
def test_bedrock_api_key_different_region(self, mock_completion):
|
|
"""Test Bedrock API key authentication with different region."""
|
|
mock_completion.return_value = MagicMock()
|
|
mock_completion.return_value.choices = [MagicMock()]
|
|
mock_completion.return_value.choices[0].message.content = "Test response"
|
|
|
|
llm = LLM(model="bedrock/anthropic.claude-3-sonnet-20240229-v1:0")
|
|
result = llm.call("test message")
|
|
|
|
mock_completion.assert_called_once()
|
|
assert result == "Test response"
|
|
@patch.dict(os.environ, {
|
|
'AWS_BEARER_TOKEN_BEDROCK': 'test-api-key',
|
|
'AWS_DEFAULT_REGION': 'us-east-1'
|
|
})
|
|
@patch('litellm.completion')
|
|
def test_bedrock_timeout_handling(self, mock_completion):
|
|
"""Test Bedrock API timeout handling."""
|
|
mock_completion.side_effect = TimeoutError("Request timed out")
|
|
|
|
llm = LLM(model="bedrock/anthropic.claude-3-sonnet-20240229-v1:0")
|
|
with pytest.raises(TimeoutError, match="Request timed out"):
|
|
llm.call("test message")
|
|
|
|
@patch.dict(os.environ, {
|
|
'AWS_BEARER_TOKEN_BEDROCK': 'test-api-key',
|
|
'AWS_DEFAULT_REGION': 'us-east-1'
|
|
})
|
|
@patch('litellm.completion')
|
|
def test_bedrock_rate_limit_handling(self, mock_completion):
|
|
"""Test Bedrock API rate limit handling."""
|
|
mock_completion.side_effect = Exception("Rate limit exceeded")
|
|
|
|
llm = LLM(model="bedrock/anthropic.claude-3-sonnet-20240229-v1:0")
|
|
with pytest.raises(Exception, match="Rate limit exceeded"):
|
|
llm.call("test message")
|
|
|
|
@patch.dict(os.environ, {
|
|
'AWS_BEARER_TOKEN_BEDROCK': 'invalid-key',
|
|
'AWS_DEFAULT_REGION': 'us-east-1'
|
|
})
|
|
@patch('litellm.completion')
|
|
def test_bedrock_invalid_api_key(self, mock_completion):
|
|
"""Test Bedrock with invalid API key."""
|
|
mock_completion.side_effect = Exception("Invalid API key")
|
|
|
|
llm = LLM(model="bedrock/anthropic.claude-3-sonnet-20240229-v1:0")
|
|
with pytest.raises(Exception, match="Invalid API key"):
|
|
llm.call("test message")
|
|
|
|
@patch.dict(os.environ, {
|
|
'AWS_ACCESS_KEY_ID': 'test-key-id',
|
|
'AWS_SECRET_ACCESS_KEY': 'test-secret-key',
|
|
'AWS_DEFAULT_REGION': 'us-east-1'
|
|
})
|
|
@patch('litellm.completion')
|
|
def test_bedrock_connection_error(self, mock_completion):
|
|
"""Test Bedrock with connection error."""
|
|
mock_completion.side_effect = ConnectionError("Connection failed")
|
|
|
|
llm = LLM(model="bedrock/anthropic.claude-3-sonnet-20240229-v1:0")
|
|
with pytest.raises(ConnectionError, match="Connection failed"):
|
|
llm.call("test message")
|
|
|
|
@patch.dict(os.environ, {
|
|
'AWS_BEARER_TOKEN_BEDROCK': 'test-api-key',
|
|
'AWS_DEFAULT_REGION': 'us-east-1'
|
|
})
|
|
@patch('litellm.completion')
|
|
def test_bedrock_api_key_with_retry_scenario(self, mock_completion):
|
|
"""Test Bedrock API key authentication with retry scenario."""
|
|
mock_completion.side_effect = [
|
|
Exception("Temporary error"),
|
|
MagicMock(choices=[MagicMock(message=MagicMock(content="Success after retry"))])
|
|
]
|
|
|
|
llm = LLM(model="bedrock/anthropic.claude-3-sonnet-20240229-v1:0")
|
|
with pytest.raises(Exception, match="Temporary error"):
|
|
llm.call("test message")
|