From 9581d80b5fb90dec9c84397ede80ee7fd05f6f0c Mon Sep 17 00:00:00 2001
From: "Chi Song (from Dev Box)" <chisong@microsoft.com>
Date: Tue, 25 Jun 2024 14:05:01 -0700
Subject: [PATCH] test: fix serial console auth needed.

---
 lisa/sut_orchestrator/azure/features.py | 30 ++++++++++++++++++-------
 1 file changed, 22 insertions(+), 8 deletions(-)

diff --git a/lisa/sut_orchestrator/azure/features.py b/lisa/sut_orchestrator/azure/features.py
index 86b2009b79..ec39961dcf 100644
--- a/lisa/sut_orchestrator/azure/features.py
+++ b/lisa/sut_orchestrator/azure/features.py
@@ -326,7 +326,7 @@ def write(self, data: str) -> None:
             return
         except websockets.ConnectionClosed as e:  # type: ignore
             # If the connection is closed, we need to reconnect
-            self._log.debug(f"Connection closed: {e}")
+            self._log.debug(f"Connection closed on read serial console: {e}")
             self._ws = None
             self._get_connection()
             raise e
@@ -340,7 +340,7 @@ def read(self) -> str:
             return output
         except websockets.ConnectionClosed as e:  # type: ignore
             # If the connection is closed, we need to reconnect
-            self._log.debug(f"Connection closed: {e}")
+            self._log.debug(f"Connection closed on read serial console: {e}")
             self._ws = None
             self._get_connection()
             raise e
@@ -360,10 +360,18 @@ def _get_connection(self) -> Any:
             connection_str = self._get_connection_string()
 
             # create websocket connection
-            self._ws = self._get_event_loop().run_until_complete(
+            ws = self._get_event_loop().run_until_complete(
                 websockets.connect(connection_str)  # type: ignore
             )
 
+            token = self._get_access_token()
+            # add to secret in case it's echo back.
+            add_secret(token)
+            # send token to auth
+            self._get_event_loop().run_until_complete(ws.send(token))
+
+            self._ws = ws
+
         return self._ws
 
     def _write(self, cmd: str) -> None:
@@ -402,6 +410,14 @@ def _read(self) -> str:
 
         return output
 
+    def _get_access_token(self) -> str:
+        platform: AzurePlatform = self._platform  # type: ignore
+        access_token = platform.credential.get_token(
+            "https://management.core.windows.net/.default"
+        ).token
+
+        return access_token
+
     def _get_console_log(self, saved_path: Optional[Path]) -> bytes:
         platform: AzurePlatform = self._platform  # type: ignore
         return save_console_log(
@@ -414,7 +430,6 @@ def _get_console_log(self, saved_path: Optional[Path]) -> bytes:
 
     def _get_connection_string(self) -> str:
         # setup connection string
-        platform: AzurePlatform = self._platform  # type: ignore
         connection = self._serial_port_operations.connect(
             resource_group_name=self._resource_group_name,
             resource_provider_namespace=self.RESOURCE_PROVIDER_NAMESPACE,
@@ -422,11 +437,8 @@ def _get_connection_string(self) -> str:
             parent_resource=self._vm_name,
             serial_port=self._serial_port.name,
         )
-        access_token = platform.credential.get_token(
-            "https://management.core.windows.net/.default"
-        ).token
         serial_port_connection_str = (
-            f"{connection.connection_string}?authorization={access_token}"
+            f"{connection.connection_string}?authorization={self._get_access_token()}"
         )
 
         return serial_port_connection_str
@@ -476,6 +488,8 @@ def _initialize_serial_console(self, port_id: int) -> None:
                 if int(serialport.name) == port_id
             ][0]
 
+        self._log.debug(f"Serial port {port_id} is enabled: {self._serial_port}")
+
         # setup shared web socket connection variable
         self._ws = None