From 2ba89870e94f11371ea67e46703bdcca2c273cc9 Mon Sep 17 00:00:00 2001 From: Josh Usiskin <56369778+jusiskin@users.noreply.github.com> Date: Fri, 27 Oct 2023 13:42:44 -0500 Subject: [PATCH] fix!: CreateFleet API changes and robust cleanup (#35) Signed-off-by: Josh Usiskin <56369778+jusiskin@users.noreply.github.com> --- .../deadline/resources.py | 7 +- src/deadline_test_fixtures/fixtures.py | 122 ++++++++++-------- test/unit/deadline/test_resources.py | 2 + 3 files changed, 79 insertions(+), 52 deletions(-) diff --git a/src/deadline_test_fixtures/deadline/resources.py b/src/deadline_test_fixtures/deadline/resources.py index 491e3470..0e7f572a 100644 --- 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 index fe8719dc..b32cdb6c 100644 --- 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 index 34e9883e..8559ab68 100644 --- 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,