fix!: CreateFleet API changes and robust cleanup (aws-deadline#35)

Signed-off-by: Josh Usiskin <[email protected]>
jusiskin · Oct 27, 2023 · 2ba8987 · 2ba8987
1 parent 64eded7
commit 2ba8987
Show file tree

Hide file tree

Showing 3 changed files with 79 additions and 52 deletions.
diff --git a/src/deadline_test_fixtures/deadline/resources.py b/src/deadline_test_fixtures/deadline/resources.py
@@ -60,9 +60,9 @@ def create(
         client: DeadlineClient,
         display_name: str,
         farm: Farm,
+        job_run_as_user: JobRunAsUser,
         role_arn: str | None = None,
         job_attachments: JobAttachmentSettings | None = None,
-        job_run_as_user: JobRunAsUser,
         raw_kwargs: dict | None = None,
     ) -> Queue:
         kwargs = clean_kwargs(
@@ -111,6 +111,8 @@ def create(
         display_name: str,
         farm: Farm,
         configuration: dict,
+        max_worker_count: int,
+        min_worker_count: int | None = None,
         role_arn: str | None = None,
         raw_kwargs: dict | None = None,
     ) -> Fleet:
@@ -120,9 +122,12 @@ def create(
                 "displayName": display_name,
                 "roleArn": role_arn,
                 "configuration": configuration,
+                "maxWorkerCount": max_worker_count,
                 **(raw_kwargs or {}),
             }
         )
+        if min_worker_count is not None:
+            kwargs["minWorkerCount"] = min_worker_count
         response = call_api(
             fn=lambda: client.create_fleet(**kwargs),
             description=f"Create fleet {display_name} in farm {farm.id}",

diff --git a/src/deadline_test_fixtures/fixtures.py b/src/deadline_test_fixtures/fixtures.py
@@ -11,8 +11,9 @@
 import posixpath
 import pytest
 import tempfile
+from contextlib import ExitStack, contextmanager
 from dataclasses import InitVar, dataclass, field, fields, MISSING
-from typing import Any, Generator
+from typing import Any, Generator, TypeVar
 
 from .deadline.client import DeadlineClient
 from .deadline.resources import (
@@ -280,58 +281,77 @@ def deadline_resources(
     else:
         LOG.info("Deploying Deadline resources")
         bootstrap_resources: BootstrapResources = request.getfixturevalue("bootstrap_resources")
-        farm = Farm.create(
-            client=deadline_client,
-            display_name="test-scaffolding-farm",
-        )
-        queue = Queue.create(
-            client=deadline_client,
-            display_name="test-scaffolding-queue",
-            farm=farm,
-            job_attachments=bootstrap_resources.job_attachments,
-            role_arn=bootstrap_resources.session_role_arn,
-            job_run_as_user=bootstrap_resources.job_run_as_user,
-        )
-        fleet = Fleet.create(
-            client=deadline_client,
-            display_name="test-scaffolding-fleet",
-            farm=farm,
-            configuration={
-                "customerManaged": {
-                    "autoScalingConfiguration": {
-                        "mode": "NO_SCALING",
-                        "maxFleetSize": 1,
-                    },
-                    "workerRequirements": {
-                        "vCpuCount": {"min": 1},
-                        "memoryMiB": {"min": 1024},
-                        "osFamily": "linux",
-                        "cpuArchitectureType": "x86_64",
-                    },
-                },
-            },
-            role_arn=bootstrap_resources.worker_role_arn,
-        )
-        qfa = QueueFleetAssociation.create(
-            client=deadline_client,
-            farm=farm,
-            queue=queue,
-            fleet=fleet,
-        )
 
-        yield DeadlineResources(
-            farm_id=farm.id,
-            queue_id=queue.id,
-            fleet_id=fleet.id,
-            job_attachments_bucket=bootstrap_resources.job_attachments.bucket_name
-            if bootstrap_resources.job_attachments
-            else None,
-        )
+        # Define a context manager for robust cleanup of resources
+        T = TypeVar("T", Farm, Fleet, Queue, QueueFleetAssociation)
+
+        @contextmanager
+        def deletable(resource: T) -> Generator[T, None, None]:
+            yield resource
+            resource.delete(client=deadline_client)
+
+        with ExitStack() as context_stack:
+            farm = context_stack.enter_context(
+                deletable(
+                    Farm.create(
+                        client=deadline_client,
+                        display_name="test-scaffolding-farm",
+                    )
+                )
+            )
+            queue = context_stack.enter_context(
+                deletable(
+                    Queue.create(
+                        client=deadline_client,
+                        display_name="test-scaffolding-queue",
+                        farm=farm,
+                        job_attachments=bootstrap_resources.job_attachments,
+                        role_arn=bootstrap_resources.session_role_arn,
+                        job_run_as_user=bootstrap_resources.job_run_as_user,
+                    )
+                )
+            )
+            fleet = context_stack.enter_context(
+                deletable(
+                    Fleet.create(
+                        client=deadline_client,
+                        display_name="test-scaffolding-fleet",
+                        farm=farm,
+                        configuration={
+                            "customerManaged": {
+                                "mode": "NO_SCALING",
+                                "workerRequirements": {
+                                    "vCpuCount": {"min": 1},
+                                    "memoryMiB": {"min": 1024},
+                                    "osFamily": "linux",
+                                    "cpuArchitectureType": "x86_64",
+                                },
+                            },
+                        },
+                        max_worker_count=1,
+                        role_arn=bootstrap_resources.worker_role_arn,
+                    )
+                )
+            )
+            context_stack.enter_context(
+                deletable(
+                    QueueFleetAssociation.create(
+                        client=deadline_client,
+                        farm=farm,
+                        queue=queue,
+                        fleet=fleet,
+                    )
+                )
+            )
 
-        qfa.delete(client=deadline_client)
-        fleet.delete(client=deadline_client)
-        queue.delete(client=deadline_client)
-        farm.delete(client=deadline_client)
+            yield DeadlineResources(
+                farm_id=farm.id,
+                queue_id=queue.id,
+                fleet_id=fleet.id,
+                job_attachments_bucket=bootstrap_resources.job_attachments.bucket_name
+                if bootstrap_resources.job_attachments
+                else None,
+            )
 
 
 @pytest.fixture(scope="session")

diff --git a/test/unit/deadline/test_resources.py b/test/unit/deadline/test_resources.py
@@ -149,6 +149,7 @@ def test_create(self, farm: Farm) -> None:
                 farm=farm,
                 configuration=configuration,
                 role_arn=role_arn,
+                max_worker_count=1,
             )
 
         # THEN
@@ -158,6 +159,7 @@ def test_create(self, farm: Farm) -> None:
             displayName=display_name,
             roleArn=role_arn,
             configuration=configuration,
+            maxWorkerCount=1,
         )
         mock_wait_for_desired_status.assert_called_once_with(
             client=mock_client,