mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-01-10 16:48:30 +00:00
3eea890409
- Remove unused pytest import from test_bedrock_authentication.py - Fixes F401 lint error that was causing CI failure Co-Authored-By: Jo\u00E3o <joao@crewai.com>
109 lines
4.3 KiB
Python
109 lines
4.3 KiB
Python
import os
|
|
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"
|