fix: Use status instead of targetStatus in UpdateWorker.

Signed-off-by: Karthik Bekal Pattathana <pattatk@dev-dsk-pattatk-2b-68804932.us-west-2.amazon.com>
aws-deadline · Sep 13, 2023 · 8abb05f · 8abb05f
1 parent bf746f2
commit 8abb05f
Show file tree

Hide file tree

Showing 7 changed files with 20 additions and 17 deletions.
diff --git a/src/deadline_worker_agent/aws/deadline/__init__.py b/src/deadline_worker_agent/aws/deadline/__init__.py
@@ -507,7 +507,7 @@ def update_worker(
         fleetId=config.fleet_id,
         workerId=worker_id,
         capabilities=config.capabilities.for_update_worker(),
-        targetStatus=status.value,
+        status=status.value,
     )
     if host_properties:
         request["hostProperties"] = host_properties

diff --git a/src/deadline_worker_agent/boto/shim.py b/src/deadline_worker_agent/boto/shim.py
@@ -416,7 +416,7 @@ def update_worker(
         farmId: str,
         fleetId: str,
         workerId: str,
-        targetStatus: Union[str, WorkerStatus],
+        status: Union[str, WorkerStatus],
         capabilities: dict[str, Any] | None = None,
         hostProperties: HostProperties | None = None,
     ) -> UpdateWorkerResponse:
@@ -425,7 +425,7 @@ def update_worker(
                 "farmId": farmId,
                 "fleetId": fleetId,
                 "workerId": workerId,
-                "targetStatus": targetStatus,
+                "status": status,
             }
             if capabilities is not None:
                 request["capabilities"] = capabilities

diff --git a/src/deadline_worker_agent/scheduler/scheduler.py b/src/deadline_worker_agent/scheduler/scheduler.py
@@ -416,7 +416,10 @@ def _transition_to_stopping(self, timeout: timedelta) -> None:
         given additional new tasks to work on.
         """
         request = dict[str, Any](
-            farmId=self._farm_id, fleetId=self._fleet_id, workerId=self._worker_id, state="STOPPING"
+            farmId=self._farm_id,
+            fleetId=self._fleet_id,
+            workerId=self._worker_id,
+            status="STOPPING",
         )
 
         start_time = monotonic()

diff --git a/src/deadline_worker_agent/session_events.py b/src/deadline_worker_agent/session_events.py
@@ -46,7 +46,7 @@ class LoggingAllowList(TypedDict, total=False):
     "deadline.AssumeQueueRoleForWorker": {"log_request_url": True},
     "deadline.UpdateWorker": {
         "log_request_url": True,
-        "req_body_keys": ["targetStatus"],
+        "req_body_keys": ["status"],
         "res_body_keys": ["log"],
     },
     "deadline.BatchGetJobEntity": {"log_request_url": True},

diff --git a/test/unit/aws/deadline/test_update_worker.py b/test/unit/aws/deadline/test_update_worker.py
@@ -117,7 +117,7 @@ def test_success(
             fleetId=config.fleet_id,
             workerId=worker_id,
             capabilities=config.capabilities.for_update_worker(),
-            targetStatus=status.value,
+            status=status.value,
             hostProperties=host_properties,
         )
     else:
@@ -126,7 +126,7 @@ def test_success(
             fleetId=config.fleet_id,
             workerId=worker_id,
             capabilities=config.capabilities.for_update_worker(),
-            targetStatus=status.value,
+            status=status.value,
         )
 
 
@@ -180,23 +180,23 @@ def test_updates_to_stopped_if_required(
                 fleetId=config.fleet_id,
                 workerId=worker_id,
                 capabilities=config.capabilities.for_update_worker(),
-                targetStatus=WorkerStatus.STARTED.value,
+                status=WorkerStatus.STARTED.value,
                 hostProperties=HOST_PROPERTIES,
             ),
             call(
                 farmId=config.farm_id,
                 fleetId=config.fleet_id,
                 workerId=worker_id,
                 capabilities=config.capabilities.for_update_worker(),
-                targetStatus=WorkerStatus.STOPPED.value,
+                status=WorkerStatus.STOPPED.value,
                 hostProperties=HOST_PROPERTIES,
             ),
             call(
                 farmId=config.farm_id,
                 fleetId=config.fleet_id,
                 workerId=worker_id,
                 capabilities=config.capabilities.for_update_worker(),
-                targetStatus=WorkerStatus.STARTED.value,
+                status=WorkerStatus.STARTED.value,
                 hostProperties=HOST_PROPERTIES,
             ),
         )
@@ -300,15 +300,15 @@ def test_reraises_when_updates_to_stopped(
                 fleetId=config.fleet_id,
                 workerId=worker_id,
                 capabilities=config.capabilities.for_update_worker(),
-                targetStatus=WorkerStatus.STARTED.value,
+                status=WorkerStatus.STARTED.value,
                 hostProperties=HOST_PROPERTIES,
             ),
             call(
                 farmId=config.farm_id,
                 fleetId=config.fleet_id,
                 workerId=worker_id,
                 capabilities=config.capabilities.for_update_worker(),
-                targetStatus=WorkerStatus.STOPPED.value,
+                status=WorkerStatus.STOPPED.value,
                 hostProperties=HOST_PROPERTIES,
             ),
         )

diff --git a/test/unit/scheduler/test_scheduler.py b/test/unit/scheduler/test_scheduler.py
@@ -344,7 +344,7 @@ def test_updates_to_stopping(self, scheduler: WorkerScheduler) -> None:
 
             # THEN
             api_mock.assert_called_once()
-            assert api_mock.call_args.kwargs["state"] == "STOPPING"
+            assert api_mock.call_args.kwargs["status"] == "STOPPING"
 
     @pytest.mark.parametrize("code", ["ThrottlingException", "InternalServerException"])
     def test_retries_on_exception(self, scheduler: WorkerScheduler, code: str) -> None:

diff --git a/test/unit/test_session_events.py b/test/unit/test_session_events.py
@@ -18,7 +18,7 @@ def test_logging_allow_list():
         "deadline.AssumeQueueRoleForWorker": {"log_request_url": True},
         "deadline.UpdateWorker": {
             "log_request_url": True,
-            "req_body_keys": ["targetStatus"],
+            "req_body_keys": ["status"],
             "res_body_keys": ["log"],
         },
         "deadline.BatchGetJobEntity": {"log_request_url": True},
@@ -66,10 +66,10 @@ def test_logging_allow_list():
                     "Content-Type": "application/json",
                     "User-Agent": "Boto3/1.28.12 md/Botocore#1.31.12 ua/2.0 os/linux#5.4.247-169.350.amzn2int.x86_64 md/arch#x86_64 lang/python#3.9.17 md/pyimpl#CPython cfg/retry-mode#legacy Botocore/1.31.12",
                 },
-                "body": b'{"targetStatus": "STARTED", "capabilities": {"amounts": [{"name": "amount.worker.vcpu", "value": 8.0}, {"name": "amount.worker.memory", "value": 14987.5234375}, {"name": "amount.worker.disk.scratch", "value": 0.0}, {"name": "amount.worker.gpu", "value": 0.0}, {"name": "amount.worker.gpu.memory", "value": 0.0}], "attributes": [{"name": "attr.worker.os.family", "values": ["linux"]}, {"name": "attr.worker.cpu.arch", "values": ["x86_64"]}]}, "hostProperties": {"ipAddresses": {"ipV4Addresses": ["127.0.0.1", "0.0.0.0", "0.0.0.0"], "ipV6Addresses": ["0000:0000:0000:0000:0000:0000:00000:0000", "0000:0000:0000:0000:0000:0000:0000:0001", "0000:0000:0000:0000:0000:0000:0000:0000"]}, "hostName": "host.us-west-2.amazon.com"}}',
+                "body": b'{"status": "STARTED", "capabilities": {"amounts": [{"name": "amount.worker.vcpu", "value": 8.0}, {"name": "amount.worker.memory", "value": 14987.5234375}, {"name": "amount.worker.disk.scratch", "value": 0.0}, {"name": "amount.worker.gpu", "value": 0.0}, {"name": "amount.worker.gpu.memory", "value": 0.0}], "attributes": [{"name": "attr.worker.os.family", "values": ["linux"]}, {"name": "attr.worker.cpu.arch", "values": ["x86_64"]}]}, "hostProperties": {"ipAddresses": {"ipV4Addresses": ["127.0.0.1", "0.0.0.0", "0.0.0.0"], "ipV6Addresses": ["0000:0000:0000:0000:0000:0000:00000:0000", "0000:0000:0000:0000:0000:0000:0000:0001", "0000:0000:0000:0000:0000:0000:0000:0000"]}, "hostName": "host.us-west-2.amazon.com"}}',
                 "url": "https://**********.execute-api.us-west-2.amazonaws.com/2020-08-21/farms/farm-0000000000000000000000000000000/fleets/fleet-0000000000000000000000000000000/workers/worker-0000000000000000000000000000000",
             },
-            "{'log_type': 'boto_request', 'operation': 'deadline.UpdateWorker', 'params': {'targetStatus': 'STARTED'}, 'request_url': 'https://**********.execute-api.us-west-2.amazonaws.com/2020-08-21/farms/farm-0000000000000000000000000000000/fleets/fleet-0000000000000000000000000000000/workers/worker-0000000000000000000000000000000'}",
+            "{'log_type': 'boto_request', 'operation': 'deadline.UpdateWorker', 'params': {'status': 'STARTED'}, 'request_url': 'https://**********.execute-api.us-west-2.amazonaws.com/2020-08-21/farms/farm-0000000000000000000000000000000/fleets/fleet-0000000000000000000000000000000/workers/worker-0000000000000000000000000000000'}",
             id="UpdateWorkerBeforeCallTest",
         ),
         pytest.param(
@@ -172,7 +172,7 @@ def test_log_before_call(
                     "RequestId": "abc878ee-32b5-44d4-885f-29071648328c",
                     "HTTPStatusCode": 200,
                 },
-                "targetStatus": "STOPPED",
+                "status": "STOPPED",
             },
             "{'log_type': 'boto_response', 'operation': 'deadline.UpdateWorker', 'status_code': 200, 'params': {'log': None}}",
             id="UpdateWorkerAfterCallTest",