From 527511281ed1dff9cbec980e66a53c5172f82090 Mon Sep 17 00:00:00 2001
From: Charles Moore <moorec@amazon.com>
Date: Tue, 16 Jul 2024 19:16:52 +0000
Subject: [PATCH] feat: Add stop/start service method

Signed-off-by: Charles Moore <moorec@amazon.com>
---
 src/deadline_test_fixtures/deadline/worker.py | 43 +++++++++++++++++--
 1 file changed, 40 insertions(+), 3 deletions(-)

diff --git a/src/deadline_test_fixtures/deadline/worker.py b/src/deadline_test_fixtures/deadline/worker.py
index 4a62034..91f6639 100644
--- a/src/deadline_test_fixtures/deadline/worker.py
+++ b/src/deadline_test_fixtures/deadline/worker.py
@@ -155,6 +155,14 @@ def configure_worker_command(
     ) -> str:  # pragma: no cover
         raise NotImplementedError("'configure_worker_command' was not implemented.")
 
+    @abc.abstractmethod
+    def start_worker_service(self) -> None:  # pragma: no cover
+        raise NotImplementedError("'_start_worker_service' was not implemented.")
+
+    @abc.abstractmethod
+    def stop_worker_service(self) -> None:  # pragma: no cover
+        raise NotImplementedError("'_stop_worker_service' was not implemented.")
+
     @abc.abstractmethod
     def get_worker_id(self) -> str:
         raise NotImplementedError("'get_worker_id' was not implemented.")
@@ -422,8 +430,8 @@ def _start_worker_agent(self) -> None:
         cmd_result = self.send_command(
             " ; ".join(
                 [
-                    "echo Waiting 20s for the agent service to get started",
-                    "sleep 20",
+                    "echo Waiting 30s for the agent service to get started",
+                    "sleep 30",
                     "echo 'Running running Get-Process to check if the agent is running'",
                     "IF(Get-Process pythonservice){echo 'service is running'}ELSE{exit 1}",
                 ]
@@ -432,6 +440,9 @@ def _start_worker_agent(self) -> None:
         assert cmd_result.exit_code == 0, f"Failed to start Worker agent: {cmd_result}"
         LOG.info("Successfully started Worker agent")
 
+        # Wait 20 seconds for the worker to start before attempting to get ID
+        time.sleep(20)
+
         self.worker_id = self.get_worker_id()
 
     def configure_worker_command(self, *, config: DeadlineWorkerConfiguration) -> str:
@@ -508,6 +519,19 @@ def userdata(self, s3_files) -> str:
 
         return userdata
 
+    def start_worker_service(self):
+        LOG.info("Sending command to start the Worker Agent service")
+
+        cmd_result = self.send_command('Start-Service -Name "DeadlineWorker"')
+
+        assert cmd_result.exit_code == 0, f"Failed to start Worker Agent service: : {cmd_result}"
+
+    def stop_worker_service(self):
+        LOG.info("Sending command to stop the Worker Agent service")
+        cmd_result = self.send_command('Stop-Service -Name "DeadlineWorker"')
+
+        assert cmd_result.exit_code == 0, f"Failed to stop Worker Agent service: : {cmd_result}"
+
     def ami_ssm_param_name(self) -> str:
         # Grab the latest Windows Server 2022 AMI
         # https://aws.amazon.com/blogs/mt/query-for-the-latest-windows-ami-using-systems-manager-parameter-store/
@@ -551,7 +575,7 @@ def _start_worker_agent(self) -> None:
         LOG.info(f"Sending SSM command to configure Worker agent on instance {self.instance_id}")
 
         cmd_result = self.send_command(
-            f"{self.configure_worker_command(config=self.configuration)}"
+            f"cd /home/{self.configuration.agent_user}; . .venv/bin/activate; AWS_DEFAULT_REGION={self.configuration.region} {self.configure_worker_command(config=self.configuration)}"
         )
         assert cmd_result.exit_code == 0, f"Failed to configure Worker agent: {cmd_result}"
         LOG.info("Successfully configured Worker agent")
@@ -666,6 +690,19 @@ def userdata(self, s3_files) -> str:
 
         return userdata
 
+    def start_worker_service(self):
+        LOG.info("Sending command to start the Worker Agent service")
+
+        cmd_result = self.send_command("systemctl start deadline-worker")
+
+        assert cmd_result.exit_code == 0, f"Failed to start Worker Agent service: {cmd_result}"
+
+    def stop_worker_service(self):
+        LOG.info("Sending command to stop the Worker Agent service")
+        cmd_result = self.send_command("systemctl stop deadline-worker")
+
+        assert cmd_result.exit_code == 0, f"Failed to start Worker Agent service: {cmd_result}"
+
     def get_worker_id(self) -> str:
         cmd_result = self.send_command("cat /var/lib/deadline/worker.json  | jq -r '.worker_id'")
         assert cmd_result.exit_code == 0, f"Failed to get Worker ID: {cmd_result}"