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 + ), )