From 9a753d03d87ec89a969dab25e0c74dc15d23809c Mon Sep 17 00:00:00 2001
From: Trevor Blanchard <blantrev@amazon.com>
Date: Sat, 23 Mar 2024 18:39:24 -0500
Subject: [PATCH] fix: FleetAPI compatibility for WorkerRequirements

Signed-off-by: Trevor Blanchard <blantrev@amazon.com>
---
 src/deadline_test_fixtures/deadline/client.py | 18 ++++++++++++++++++
 src/deadline_test_fixtures/fixtures.py        | 10 ++++++----
 2 files changed, 24 insertions(+), 4 deletions(-)

diff --git a/src/deadline_test_fixtures/deadline/client.py b/src/deadline_test_fixtures/deadline/client.py
index d1d4834..b97e691 100644
--- 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
index bb882c6..776e810 100644
--- 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
+                ),
             )