Implement serial log on Azure

lisa/sut_orchestrator/azure/features.py

@@ -1,11 +1,22 @@
-from typing import Any
+from pathlib import Path
+from typing import Any, Optional
 
-from lisa.features import StartStop as BaseStartStop
+import requests
+
+from lisa import features
+from lisa.node import Node
 
 from .common import get_compute_client, get_node_context, wait_operation
 
 
-class StartStop(BaseStartStop):
+class AzureFeatureMixin:
+    def _initialize_information(self, node: Node) -> None:
+        node_context = get_node_context(node)
+        self._vm_name = node_context.vm_name
+        self._resource_group_name = node_context.resource_group_name
+
+
+class StartStop(AzureFeatureMixin, features.StartStop):
     def _stop(self, wait: bool = True) -> Any:
         return self._execute(wait, "begin_stop")
 
@@ -16,9 +27,8 @@ class StartStop(BaseStartStop):
         return self._execute(wait, "begin_restart")
 
     def _initialize(self, *args: Any, **kwargs: Any) -> None:
-        node_context = get_node_context(self._node)
-        self._vm_name = node_context.vm_name
-        self._resource_group_name = node_context.resource_group_name
+        super()._initialize(*args, **kwargs)
+        self._initialize_information(self._node)
 
     def _execute(self, wait: bool, operator: str) -> Any:
         compute_client = get_compute_client(self._platform)
@@ -29,3 +39,28 @@ class StartStop(BaseStartStop):
         if wait:
             result = wait_operation(result)
         return result
+
+
+class SerialConsole(AzureFeatureMixin, features.SerialConsole):
+    def _initialize(self, *args: Any, **kwargs: Any) -> None:
+        super()._initialize(*args, **kwargs)
+        self._initialize_information(self._node)
+
+    def _get_console_log(self, saved_path: Optional[Path]) -> bytes:
+        compute_client = get_compute_client(self._platform)
+        diagnostic_data = (
+            compute_client.virtual_machines.retrieve_boot_diagnostics_data(
+                resource_group_name=self._resource_group_name, vm_name=self._vm_name
+            )
+        )
+        if saved_path:
+            screenshot_name = saved_path.joinpath("serial_console.bmp")
+            screenshot_response = requests.get(
+                diagnostic_data.console_screenshot_blob_uri
+            )
+            with open(screenshot_name, mode="wb") as f:
+                f.write(screenshot_response.content)
+
+        log_response = requests.get(diagnostic_data.serial_console_log_blob_uri)
+
+        return log_response.content

lisa/sut_orchestrator/azure/platform_.py

@@ -322,7 +322,7 @@ class AzurePlatform(Platform):
 
     @classmethod
     def supported_features(cls) -> List[Type[Feature]]:
-        return [features.StartStop]
+        return [features.StartStop, features.SerialConsole]
 
     def _prepare_environment(  # noqa: C901
         self, environment: Environment, log: Logger
@@ -987,7 +987,9 @@ class AzurePlatform(Platform):
 
         # all node support start/stop
         node_space.features = search_space.SetSpace[str](is_allow_set=True)
-        node_space.features.add(features.StartStop.name())
+        node_space.features.update(
+            [features.StartStop.name(), features.SerialConsole.name()]
+        )
 
         return node_space