fix: FleetAPI compatibility for WorkerRequirements

Signed-off-by: Trevor Blanchard <[email protected]>
aws-deadline · Mar 24, 2024 · 9a753d0 · 9a753d0
1 parent a1fa350
commit 9a753d0
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 4 deletions.
diff --git a/src/deadline_test_fixtures/deadline/client.py b/src/deadline_test_fixtures/deadline/client.py
@@ -36,6 +36,24 @@ def create_fleet(self, *args, **kwargs) -> Any:
             and "workeRoleArn" in create_fleet_input_members
         ):
             kwargs["workerRoleArn"] = kwargs.pop("roleArn")
+
+        if "configuration" in create_fleet_input_members:
+            customer_managed_members = (
+                create_fleet_input_members["configuration"].members["customerManaged"].members
+            )
+
+            if (
+                "workerCapabilities" not in customer_managed_members
+                and "workerRequirements" in customer_managed_members
+            ):
+                if (
+                    "customerManaged" in kwargs["configuration"]
+                    and "workerCapabilities" in kwargs["configuration"]["customerManaged"]
+                ):
+                    kwargs["configuration"]["customerManaged"]["workerRequirements"] = kwargs[
+                        "configuration"
+                    ]["customerManaged"].pop("workerCapabilities")
+
         return self._real_client.create_fleet(*args, **kwargs)
 
     def get_fleet(self, *args, **kwargs) -> Any:

diff --git a/src/deadline_test_fixtures/fixtures.py b/src/deadline_test_fixtures/fixtures.py
@@ -324,7 +324,7 @@ def deletable(resource: T) -> Generator[T, None, None]:
                         configuration={
                             "customerManaged": {
                                 "mode": "NO_SCALING",
-                                "workerRequirements": {
+                                "workerCapabilities": {
                                     "vCpuCount": {"min": 1},
                                     "memoryMiB": {"min": 1024},
                                     "osFamily": "linux",
@@ -352,9 +352,11 @@ def deletable(resource: T) -> Generator[T, None, None]:
                 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,
+                job_attachments_bucket=(
+                    bootstrap_resources.job_attachments.bucket_name
+                    if bootstrap_resources.job_attachments
+                    else None
+                ),
             )