fix(job-attachment): Refactor bucket creation and ensure complete res…

…ource deletion (aws-deadline#10)

Signed-off-by: Gahyun Suh <gahyusuh@amazon.com>

src/deadline_test_fixtures/cloudformation/job_attachments_bootstrap_stack.py

@@ -6,12 +6,10 @@
     BucketPolicy,
     CfnStack,
 )
-from .util import create_secure_bucket
 
 
 class JobAttachmentsBootstrapStack(CfnStack):  # pragma: no cover
     bucket: Bucket
-    log_bucket: Bucket
     bucket_policy: BucketPolicy
 
     def __init__(
@@ -23,13 +21,37 @@ def __init__(
     ) -> None:
         super().__init__(name=name, description=description)
 
-        self.bucket, self.log_bucket, self.bucket_policy = create_secure_bucket(
-            self,
-            "JobAttachmentBucket",
-            bucket_kwargs={
-                "bucket_name": bucket_name,
+        self.bucket = Bucket(
+            stack=self,
+            logical_name="JobAttachmentBucket",
+            bucket_name=bucket_name,
+            versioning=False,
+            encryption={
+                "ServerSideEncryptionConfiguration": [
+                    {"ServerSideEncryptionByDefault": {"SSEAlgorithm": "AES256"}},
+                ],
             },
-            log_bucket_kwargs={
-                "bucket_name": f"{bucket_name}-logs",
+            update_replace_policy="Delete",
+            deletion_policy="Delete",
+        )
+
+        self.bucket_policy = BucketPolicy(
+            stack=self,
+            logical_name="JobAttachmentBucketPolicy",
+            bucket=self.bucket,
+            policy_document={
+                "Version": "2012-10-17",
+                "Statement": [
+                    {
+                        "Action": "s3:*",
+                        "Effect": "Deny",
+                        "Principal": "*",
+                        "Resource": [
+                            self.bucket.arn,
+                            self.bucket.arn_for_objects(),
+                        ],
+                        "Condition": {"Bool": {"aws:SecureTransport": "false"}},
+                    },
+                ],
             },
         )

src/deadline_test_fixtures/fixtures.py

@@ -499,7 +499,7 @@ def deploy_job_attachment_resources() -> Generator[JobAttachmentManager, None, N
         JobAttachmentManager: Class to manage Job Attachments resources
     """
     manager = JobAttachmentManager(
-        s3_resource=boto3.resource("s3"),
+        s3_client=boto3.client("s3"),
         cfn_client=boto3.client("cloudformation"),
         deadline_client=DeadlineClient(boto3.client("deadline")),
         account_id=os.environ["SERVICE_ACCOUNT_ID"],

src/deadline_test_fixtures/job_attachment_manager.py

@@ -2,7 +2,6 @@
 from __future__ import annotations
 
 from dataclasses import InitVar, dataclass, field
-from typing import Any
 
 from botocore.client import BaseClient
 from botocore.exceptions import ClientError, WaiterError
@@ -21,29 +20,27 @@ class JobAttachmentManager:
     Responsible for setting up and tearing down job attachment test resources
     """
 
-    s3_resource: Any
+    s3_client: BaseClient
     cfn_client: BaseClient
     deadline_client: DeadlineClient
 
     stage: InitVar[str]
     account_id: InitVar[str]
 
-    bucket: Any = field(init=False)
     stack: JobAttachmentsBootstrapStack = field(init=False)
     farm: Farm | None = field(init=False, default=None)
     queue: Queue | None = field(init=False, default=None)
+    queue_with_no_settings: Queue | None = field(init=False, default=None)
 
     def __post_init__(
         self,
         stage: str,
         account_id: str,
     ):
-        self.bucket = self.s3_resource.Bucket(
-            f"job-attachment-integ-test-{stage.lower()}-{account_id}"
-        )
+        self.bucket_name = f"job-attachment-integ-test-{stage.lower()}-{account_id}"
         self.stack = JobAttachmentsBootstrapStack(
             name="JobAttachmentIntegTest",
-            bucket_name=self.bucket.name,
+            bucket_name=self.bucket_name,
         )
 
     def deploy_resources(self):
@@ -60,6 +57,11 @@ def deploy_resources(self):
                 display_name="job_attachments_test_queue",
                 farm=self.farm,
             )
+            self.queue_with_no_settings = Queue.create(
+                client=self.deadline_client,
+                display_name="job_attachments_test_no_settings_queue",
+                farm=self.farm,
+            )
             self.stack.deploy(cfn_client=self.cfn_client)
         except (ClientError, WaiterError):
             # If anything goes wrong, rollback
@@ -71,7 +73,15 @@ def empty_bucket(self):
         Empty the bucket between session runs
         """
         try:
-            self.bucket.objects.all().delete()
+            # List up all objects and their versions in the bucket
+            version_list = self.s3_client.list_object_versions(Bucket=self.bucket_name)
+            object_list = version_list.get("Versions", []) + version_list.get("DeleteMarkers", [])
+            # Delete all objects and versions
+            for obj in object_list:
+                self.s3_client.delete_object(
+                    Bucket=self.bucket_name, Key=obj["Key"], VersionId=obj.get("VersionId", None)
+                )
+
         except ClientError as e:
             if e.response["Error"]["Message"] != "The specified bucket does not exist":
                 raise
@@ -84,5 +94,7 @@ def cleanup_resources(self):
         self.stack.destroy(cfn_client=self.cfn_client)
         if self.queue:
             self.queue.delete(client=self.deadline_client)
+        if self.queue_with_no_settings:
+            self.queue_with_no_settings.delete(client=self.deadline_client)
         if self.farm:
             self.farm.delete(client=self.deadline_client)

test/unit/test_job_attachment_manager.py

@@ -38,7 +38,7 @@ def job_attachment_manager(
     ) -> Generator[JobAttachmentManager, None, None]:
         with mock_s3():
             yield JobAttachmentManager(
-                s3_resource=boto3.resource("s3"),
+                s3_client=boto3.client("s3"),
                 cfn_client=MagicMock(),
                 deadline_client=DeadlineClient(MagicMock()),
                 stage="test",
@@ -61,7 +61,7 @@ def test_deploys_all_resources(
 
             # THEN
             mock_farm_cls.create.assert_called_once()
-            mock_queue_cls.create.assert_called_once()
+            mock_queue_cls.create.call_count == 2
             mock_stack.deploy.assert_called_once()
 
         @pytest.mark.parametrize(
@@ -109,7 +109,7 @@ def test_cleans_up_when_error_is_raised(
     class TestEmptyBucket:
         def test_deletes_all_objects(self, job_attachment_manager: JobAttachmentManager):
             # GIVEN
-            bucket = job_attachment_manager.bucket
+            bucket = boto3.resource("s3").Bucket(job_attachment_manager.bucket_name)
             bucket.create()
             bucket.put_object(Key="test-object", Body="Hello world".encode())
             bucket.put_object(Key="test-object-2", Body="Hello world 2".encode())
@@ -152,17 +152,17 @@ def test_raises_any_other_error(
             # GIVEN
             exc = ClientError({"Error": {"Message": "test"}}, "test-operation")
             with (
-                patch.object(job_attachment_manager, "bucket") as mock_bucket,
+                patch.object(job_attachment_manager, "s3_client") as mock_s3_client,
                 pytest.raises(ClientError) as raised_exc,
             ):
-                mock_bucket.objects.all.side_effect = exc
+                mock_s3_client.list_object_versions.side_effect = exc
 
                 # WHEN
                 job_attachment_manager.empty_bucket()
 
             # THEN
             assert raised_exc.value is exc
-            mock_bucket.objects.all.assert_called_once()
+            mock_s3_client.list_object_versions.assert_called_once()
 
     def test_cleanup_resources(
         self,
@@ -186,5 +186,5 @@ def test_cleanup_resources(
         # THEN
         spy_empty_bucket.assert_called_once()
         mock_stack.destroy.assert_called_once()
-        mock_queue_cls.create.return_value.delete.assert_called_once()
+        mock_queue_cls.create.return_value.delete.call_count == 2
         mock_farm_cls.create.return_value.delete.assert_called_once()