From 679c6344597352c1ae765c4bcad11931b2f5eaa5 Mon Sep 17 00:00:00 2001 From: "lentil32 (aider)" Date: Sun, 16 Mar 2025 15:38:59 +0900 Subject: [PATCH] test: Add tests for AWS credential handling with AWS_PROFILE --- tests/basic/test_aws_credentials.py | 169 ++++++++++++++++++++++++++++ 1 file changed, 169 insertions(+) create mode 100644 tests/basic/test_aws_credentials.py diff --git a/tests/basic/test_aws_credentials.py b/tests/basic/test_aws_credentials.py new file mode 100644 index 000000000..87d2b3a63 --- /dev/null +++ b/tests/basic/test_aws_credentials.py @@ -0,0 +1,169 @@ +import os +from unittest.mock import patch + +from aider.models import Model + + +class TestAWSCredentials: + """Test AWS credential handling, especially AWS_PROFILE.""" + + def test_bedrock_model_with_aws_profile(self): + """Test that Bedrock models accept AWS_PROFILE as valid authentication.""" + # Save original environment + original_env = os.environ.copy() + + try: + # Set up test environment + os.environ.clear() + os.environ["AWS_PROFILE"] = "test-profile" + + # Create a model instance + with patch("aider.llm.litellm.validate_environment") as mock_validate: + # Mock the litellm validate_environment to return missing AWS keys + mock_validate.return_value = { + "missing_keys": ["AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY"], + "keys_in_environment": False, + } + + # Test with a bedrock model + model = Model("bedrock/anthropic.claude-v2") + + # Check that the AWS keys were removed from missing_keys + assert "AWS_ACCESS_KEY_ID" not in model.missing_keys + assert "AWS_SECRET_ACCESS_KEY" not in model.missing_keys + # With no missing keys, validation should pass + assert model.keys_in_environment + + finally: + # Restore original environment + os.environ.clear() + os.environ.update(original_env) + + def test_us_anthropic_model_with_aws_profile(self): + """Test that us.anthropic models accept AWS_PROFILE as valid authentication.""" + # Save original environment + original_env = os.environ.copy() + + try: + # Set up test environment + os.environ.clear() + os.environ["AWS_PROFILE"] = "test-profile" + + # Create a model instance + with patch("aider.llm.litellm.validate_environment") as mock_validate: + # Mock the litellm validate_environment to return missing AWS keys + mock_validate.return_value = { + "missing_keys": ["AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY"], + "keys_in_environment": False, + } + + # Test with a us.anthropic model + model = Model("us.anthropic.claude-3-7-sonnet-20250219-v1:0") + + # Check that the AWS keys were removed from missing_keys + assert "AWS_ACCESS_KEY_ID" not in model.missing_keys + assert "AWS_SECRET_ACCESS_KEY" not in model.missing_keys + # With no missing keys, validation should pass + assert model.keys_in_environment + + finally: + # Restore original environment + os.environ.clear() + os.environ.update(original_env) + + def test_non_bedrock_model_with_aws_profile(self): + """Test that non-Bedrock models do not accept AWS_PROFILE for AWS credentials.""" + # Save original environment + original_env = os.environ.copy() + + try: + # Set up test environment + os.environ.clear() + os.environ["AWS_PROFILE"] = "test-profile" + + # Create a model instance + with patch("aider.llm.litellm.validate_environment") as mock_validate: + # Mock the litellm validate_environment to return missing AWS keys + mock_validate.return_value = { + "missing_keys": ["AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY"], + "keys_in_environment": False, + } + + # Test with a non-Bedrock model + model = Model("gpt-4") + + # For non-Bedrock models, AWS credential keys should remain in missing_keys + assert "AWS_ACCESS_KEY_ID" in model.missing_keys + assert "AWS_SECRET_ACCESS_KEY" in model.missing_keys + # With missing keys, validation should fail + assert not model.keys_in_environment + + finally: + # Restore original environment + os.environ.clear() + os.environ.update(original_env) + + def test_bedrock_model_without_aws_profile(self): + """Test that Bedrock models require credentials when AWS_PROFILE is not set.""" + # Save original environment + original_env = os.environ.copy() + + try: + # Set up test environment + os.environ.clear() + + # Create a model instance + with patch("aider.llm.litellm.validate_environment") as mock_validate: + # Mock the litellm validate_environment to return missing AWS keys + mock_validate.return_value = { + "missing_keys": ["AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY"], + "keys_in_environment": False, + } + + # Test with a bedrock model without AWS_PROFILE + model = Model("bedrock/anthropic.claude-v2") + + # Without AWS_PROFILE, AWS credential keys should remain in missing_keys + assert "AWS_ACCESS_KEY_ID" in model.missing_keys + assert "AWS_SECRET_ACCESS_KEY" in model.missing_keys + # With missing keys, validation should fail + assert not model.keys_in_environment + + finally: + # Restore original environment + os.environ.clear() + os.environ.update(original_env) + + def test_mixed_missing_keys_with_aws_profile(self): + """Test that only AWS credential keys are affected by AWS_PROFILE.""" + # Save original environment + original_env = os.environ.copy() + + try: + # Set up test environment + os.environ.clear() + os.environ["AWS_PROFILE"] = "test-profile" + + # Create a model instance + with patch("aider.llm.litellm.validate_environment") as mock_validate: + # Mock the litellm validate_environment to return missing AWS keys and another key + mock_validate.return_value = { + "missing_keys": ["AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY", "ANOTHER_KEY"], + "keys_in_environment": False, + } + + # Test with a bedrock model + model = Model("bedrock/anthropic.claude-v2") + + # AWS credential keys should be removed from missing_keys + assert "AWS_ACCESS_KEY_ID" not in model.missing_keys + assert "AWS_SECRET_ACCESS_KEY" not in model.missing_keys + # But other keys should remain + assert "ANOTHER_KEY" in model.missing_keys + # With other missing keys, validation should still fail + assert not model.keys_in_environment + + finally: + # Restore original environment + os.environ.clear() + os.environ.update(original_env)