feat: Add support to fetch and match worker logs

Author: leongdl

src/deadline_test_fixtures/deadline/__init__.py

@@ -23,6 +23,7 @@
     Task,
     TaskParameterValue,
     TaskStatus,
+    WorkerLog,
 )
 from .worker import (
     CommandResult,
@@ -70,4 +71,5 @@
     "TaskStatus",
     "WindowsInstanceBuildWorker",
     "WindowsInstanceWorkerBase",
+    "WorkerLog",
 ]

src/deadline_test_fixtures/deadline/resources.py

@@ -1288,8 +1288,7 @@ def assert_log_does_not_contain(
 
 
 @dataclass
-class SessionLog:
-    session_id: str
+class CloudWatchLogs:
     logs: list[CloudWatchLogEvent]
 
     def assert_pattern_in_log(
@@ -1365,3 +1364,13 @@ def from_api_response(response: dict) -> CloudWatchLogEvent:
             message=response["message"],
             timestamp=response["timestamp"],
         )
+
+
+@dataclass
+class SessionLog(CloudWatchLogs):
+    session_id: str
+
+
+@dataclass
+class WorkerLog(CloudWatchLogs):
+    worker_id: str

src/deadline_test_fixtures/deadline/worker.py

@@ -16,15 +16,18 @@
 import tempfile
 import time
 from dataclasses import dataclass, field, InitVar, replace
-from typing import Any, ClassVar, Optional, cast
+from typing import TYPE_CHECKING, Any, ClassVar, Optional, cast
 
 from ..models import (
     PipInstall,
     PosixSessionUser,
 )
-from .resources import Fleet
+from .resources import CloudWatchLogEvent, Fleet, WorkerLog
 from ..util import call_api, wait_for
 
+if TYPE_CHECKING:
+    from botocore.paginate import PageIterator, Paginator
+
 LOG = logging.getLogger(__name__)
 
 DOCKER_CONTEXT_DIR = os.path.join(os.path.dirname(__file__), "..", "containers", "worker")
@@ -48,6 +51,15 @@ def get_worker_id(self) -> str:
         pass
 
 
+@dataclass(frozen=True)
+class WorkerLogConfig:
+    cloudwatch_log_group: str
+    """The name of the CloudWatch Log Group that the Agent log should be streamed to"""
+
+    cloudwatch_log_stream: str
+    """The name of the CloudWatch Log Stream that the Agent log should be streamed to"""
+
+
 @dataclass(frozen=True)
 class CommandResult:  # pragma: no cover
     exit_code: int
@@ -252,6 +264,46 @@ def set_stopped_status(self):
             LOG.exception(f"Failed to update worker status: {error}")
             raise
 
+    def get_worker_logs(self) -> Optional[WorkerLogConfig]:
+        """Get the log group and log stream for the worker. Retain the API structure"""
+        response = self.deadline_client.get_worker(
+            farmId=self.configuration.farm_id,
+            fleetId=self.configuration.fleet.id,
+            workerId=self.worker_id,
+        )
+        if log_config := response["log"]:
+            LOG.info(f"LogGroup structure {log_config}")
+            if log_config_options := log_config.get("options"):
+                log_group_name = log_config_options.get("logGroupName")
+                log_stream_name = log_config_options.get("logStreamName")
+                if log_group_name and log_stream_name:
+                    return WorkerLogConfig(
+                        cloudwatch_log_group=log_group_name, cloudwatch_log_stream=log_stream_name
+                    )
+        # Default, no log group yet.
+        return None
+
+    def get_all_worker_logs(self, *, logs_client: botocore.client.BaseClient) -> WorkerLog:
+        # Get the worker log group and stream from the service.
+        log_config = self.get_worker_logs()
+        if not log_config:
+            return WorkerLog(worker_id=self.worker_id, logs=[])  # type: ignore[arg-type]
+
+        filter_log_events_paginator: Paginator = logs_client.get_paginator("filter_log_events")
+        filter_log_events_pages: PageIterator = call_api(
+            description=f"Fetching log events for session {self.worker_id} in log group {log_config.cloudwatch_log_group}",
+            fn=lambda: filter_log_events_paginator.paginate(
+                logGroupName=log_config.cloudwatch_log_group,
+                logStreamNames=[log_config.cloudwatch_log_stream],
+            ),
+        )
+        log_events = filter_log_events_pages.build_full_result()
+        log_events = [CloudWatchLogEvent.from_api_response(e) for e in log_events["events"]]
+        # For debugging test cases.
+        # LOG.info(log_events)
+
+        return WorkerLog(worker_id=self.worker_id, logs=log_events)  # type: ignore[arg-type]
+
     def send_command(self, command: str) -> CommandResult:
         """Send a command via SSM to a shell on a launched EC2 instance. Once the command has fully
         finished the result of the invocation is returned.