Skip to content

Commit

Permalink
add unit tests for s3clientmanager.get_bucket_endpoint
Browse files Browse the repository at this point in the history
  • Loading branch information
@mandar242
mandar242 committed Mar 5, 2025
1 parent 7b9bfac commit 4354b31
Showing 1 changed file with 76 additions and 45 deletions.
121 changes: 76 additions & 45 deletions tests/unit/plugins/connection/aws_ssm/test_aws_ssm.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -312,11 +312,16 @@ def test_generate_commands(self, is_windows):
assert isinstance(test_command_generation, tuple)
assert isinstance(test_command_generation[0], list)
assert isinstance(test_command_generation[0][0], dict)

class TestS3ClientManager:
"""
Tests for the S3ClientManager class
"""
@patch("ansible_collections.community.aws.plugins.connection.aws_ssm.S3ClientManager.get_boto_client", return_value="mocked_s3_client")

@patch(
"ansible_collections.community.aws.plugins.connection.aws_ssm.S3ClientManager.get_boto_client",
return_value="mocked_s3_client",
)
def test_initialize_s3_client(self, mock_get_boto_client):
"""
Test initialize_s3_client()
Expand All @@ -328,9 +333,7 @@ def test_initialize_s3_client(self, mock_get_boto_client):
s3_client_manager = S3ClientManager(connection=conn)

s3_client_manager.initialize_s3_client(
region_name="us-east-2",
endpoint_url="https://mock-endpoint",
profile_name="test-profile"
region_name="us-east-2", endpoint_url="https://mock-endpoint", profile_name="test-profile"
)

assert mock_get_boto_client.call_count == 1
Expand All @@ -345,6 +348,72 @@ def test_initialize_s3_client(self, mock_get_boto_client):
assert s3_client_manager._s3_client is not None
assert s3_client_manager._s3_client == "mocked_s3_client"

def test_get_bucket_endpoint(self):
pc = PlayContext()
new_stdin = StringIO()
conn = connection_loader.get("community.aws.aws_ssm", pc, new_stdin)

# Mock get_option to return expected values
def mock_get_option(key):
options = {
"region": "us-east-2",
"profile": "test-profile",
"bucket_name": "my-bucket",
"bucket_endpoint_url": None,
}
return options.get(key)

conn.get_option = MagicMock(side_effect=mock_get_option)

s3_manager = S3ClientManager(connection=conn)

mock_tmp_client = MagicMock()
mock_tmp_client.head_bucket.return_value = {
"ResponseMetadata": {"HTTPHeaders": {"x-amz-bucket-region": "us-east-2"}}
}

mock_region_client = MagicMock()
mock_region_client.meta.endpoint_url = "https://s3.us-east-2.amazonaws.com"
mock_region_client.meta.region_name = "us-east-2"

s3_manager.get_boto_client = MagicMock(side_effect=[mock_tmp_client, mock_region_client])

endpoint, region = s3_manager.get_bucket_endpoint()

assert endpoint == "https://s3.us-east-2.amazonaws.com"
assert region == "us-east-2"

@patch("boto3.session.Session")
def test_get_boto_client(self, mock_session_cls):
pc = PlayContext()
new_stdin = StringIO()
conn = connection_loader.get("community.aws.aws_ssm", pc, new_stdin)

# Mock get_option to return expected values
def mock_get_option(key):
options = {
"access_key_id": "dummy_key",
"secret_access_key": "dummy_secret",
"session_token": "dummy_token",
"s3_addressing_style": "path",
}
return options.get(key, None)

conn.get_option = MagicMock(side_effect=mock_get_option)

s3_manager = S3ClientManager(connection=conn)

mock_session = MagicMock()
mock_client = "mock_client"
mock_session.client.return_value = mock_client
mock_session_cls.return_value = mock_session

client = s3_manager.get_boto_client(
service="s3", region_name="us-east-2", profile_name="test-profile", endpoint_url="http://example.com"
)

assert client == mock_client

def test_get_url_no_extra_args(self):
"""
Test get_url() without extra_args
Expand Down Expand Up @@ -393,7 +462,7 @@ def test_get_url_extra_args(self):
bucket_name="test_bucket",
out_path="test/path",
http_method="PUT",
extra_args=extra_args
extra_args=extra_args,
)

expected_params = {"Bucket": "test_bucket", "Key": "test/path"}
Expand All @@ -404,7 +473,6 @@ def test_get_url_extra_args(self):
)
assert result == "http://test-url-extra"


def test_generate_encryption_settings(self):
"""
Test generate_encryption_settings()
Expand All @@ -428,50 +496,13 @@ def mock_get_option(key):
s3_manager = S3ClientManager(connection=conn)
put_args, put_headers = s3_manager.generate_encryption_settings()

expected_put_args = {
"ServerSideEncryption": "aws:kms",
"SSEKMSKeyId": "my-kms-key-id"
}
expected_put_args = {"ServerSideEncryption": "aws:kms", "SSEKMSKeyId": "my-kms-key-id"}
expected_put_headers = {
"x-amz-server-side-encryption": "aws:kms",
"x-amz-server-side-encryption-aws-kms-key-id": "my-kms-key-id"
"x-amz-server-side-encryption-aws-kms-key-id": "my-kms-key-id",
}

assert put_args == expected_put_args
assert put_headers == expected_put_headers
conn.get_option.assert_any_call("bucket_sse_mode")
conn.get_option.assert_any_call("bucket_sse_kms_key_id")

@patch("boto3.session.Session")
def test_get_boto_client(self, mock_session_cls):
pc = PlayContext()
new_stdin = StringIO()
conn = connection_loader.get("community.aws.aws_ssm", pc, new_stdin)

# Mock get_option to return dummy credentials and settings.
def mock_get_option(key):
options = {
"access_key_id": "dummy_key",
"secret_access_key": "dummy_secret",
"session_token": "dummy_token",
"s3_addressing_style": "path",
}
return options.get(key, None)

conn.get_option = MagicMock(side_effect=mock_get_option)

s3_manager = S3ClientManager(connection=conn)

mock_session = MagicMock()
mock_client = "mock_client"
mock_session.client.return_value = mock_client
mock_session_cls.return_value = mock_session

client = s3_manager.get_boto_client(
service="s3",
region_name="us-east-2",
profile_name="test-profile",
endpoint_url="http://example.com"
)

assert client == mock_client

0 comments on commit 4354b31

Please sign in to comment.