Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fix(ingestion/aws_common): update iam role and aws access key tests to complete successfully when executed on EC2 instance #12252

Merged
merged 1 commit into from
Jan 1, 2025
Merged

fix(ingestion/aws_common): update iam role and aws access key tests to complete successfully when executed on EC2 instance #12252

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
29 changes: 21 additions & 8 deletions metadata-ingestion/tests/unit/test_aws_common.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,14 @@
)


@pytest.fixture
def mock_disable_ec2_metadata():
"""Disable EC2 metadata detection"""
with patch("requests.put") as mock_put:
mock_put.return_value.status_code = 404
yield mock_put


@pytest.fixture
def mock_aws_config():
return AwsConnectionConfig(
Expand All @@ -27,17 +35,19 @@ def mock_aws_config():


class TestAwsCommon:
def test_environment_detection_no_environment(self):
def test_environment_detection_no_environment(self, mock_disable_ec2_metadata):
"""Test environment detection when no AWS environment is present"""
with patch.dict(os.environ, {}, clear=True):
assert detect_aws_environment() == AwsEnvironment.UNKNOWN

def test_environment_detection_lambda(self):
def test_environment_detection_lambda(self, mock_disable_ec2_metadata):
"""Test Lambda environment detection"""
with patch.dict(os.environ, {"AWS_LAMBDA_FUNCTION_NAME": "test-function"}):
assert detect_aws_environment() == AwsEnvironment.LAMBDA

def test_environment_detection_lambda_cloudformation(self):
def test_environment_detection_lambda_cloudformation(
self, mock_disable_ec2_metadata
):
"""Test CloudFormation Lambda environment detection"""
with patch.dict(
os.environ,
Expand All @@ -48,7 +58,7 @@ def test_environment_detection_lambda_cloudformation(self):
):
assert detect_aws_environment() == AwsEnvironment.CLOUD_FORMATION

def test_environment_detection_eks(self):
def test_environment_detection_eks(self, mock_disable_ec2_metadata):
"""Test EKS environment detection"""
with patch.dict(
os.environ,
Expand All @@ -59,19 +69,19 @@ def test_environment_detection_eks(self):
):
assert detect_aws_environment() == AwsEnvironment.EKS

def test_environment_detection_app_runner(self):
def test_environment_detection_app_runner(self, mock_disable_ec2_metadata):
"""Test App Runner environment detection"""
with patch.dict(os.environ, {"AWS_APP_RUNNER_SERVICE_ID": "service-id"}):
assert detect_aws_environment() == AwsEnvironment.APP_RUNNER

def test_environment_detection_ecs(self):
def test_environment_detection_ecs(self, mock_disable_ec2_metadata):
"""Test ECS environment detection"""
with patch.dict(
os.environ, {"ECS_CONTAINER_METADATA_URI_V4": "http://169.254.170.2/v4"}
):
assert detect_aws_environment() == AwsEnvironment.ECS

def test_environment_detection_beanstalk(self):
def test_environment_detection_beanstalk(self, mock_disable_ec2_metadata):
"""Test Elastic Beanstalk environment detection"""
with patch.dict(os.environ, {"ELASTIC_BEANSTALK_ENVIRONMENT_NAME": "my-env"}):
assert detect_aws_environment() == AwsEnvironment.BEANSTALK
Expand Down Expand Up @@ -103,6 +113,7 @@ def test_ec2_metadata_token_failure(self, mock_put):
@patch("requests.put")
def test_is_running_on_ec2(self, mock_put, mock_get):
"""Test EC2 instance detection with IMDSv2"""
# Explicitly mock EC2 metadata responses
mock_put.return_value.status_code = 200
mock_put.return_value.text = "token123"
mock_get.return_value.status_code = 200
Expand Down Expand Up @@ -322,7 +333,9 @@ def test_aws_connection_config_validation_error(self):
),
],
)
def test_environment_detection_parametrized(self, env_vars, expected_environment):
def test_environment_detection_parametrized(
self, mock_disable_ec2_metadata, env_vars, expected_environment
):
"""Parametrized test for environment detection with different configurations"""
with patch.dict(os.environ, env_vars, clear=True):
assert detect_aws_environment() == expected_environment
Loading