Making admin commands work for serial console when the VM has a custom storage account attached (#9479)

* Making admin commands work for serial console when the VM has a custom storage account attached.

* Updating the test recordings in order for the tests to pass in the pipeline

---------

Co-authored-by: Razvan Ivan <[email protected]>

Author: Talos10

src/serial-console/HISTORY.rst

@@ -1,5 +1,9 @@
 Release History
 ===============
+1.0.0b3
+++++++
+* Fixed an issue where admin commands were not being sent when the VM was using a custom boot diagnostics storage account.
+
 1.0.0b2
 ++++++
 * Changed to 2024 API version, fixes Disable API to track "properties". Essentially return to 2018 format

src/serial-console/azext_serialconsole/custom.py

@@ -551,6 +551,9 @@ def wrapper():
         if self.load_websocket_url():
             def on_message(ws, _):
                 GV.trycount += 1
+                if GV.first_message:
+                    ws.send(self.access_token)
+                    GV.first_message = False
                 if func():
                     GV.loading = False
                     GV.terminating_app = True

src/serial-console/azext_serialconsole/tests/latest/recordings/test_check_resource_VM.yaml

@@ -15,7 +15,7 @@ interactions:
       Content-Type:
       - application/json
       User-Agent:
-      - AZURECLI/2.67.0 azsdk-python-core/1.31.0 Python/3.12.7 (Windows-11-10.0.26100-SP0)
+      - AZURECLI/2.81.0 azsdk-python-core/1.35.0 Python/3.12.3 (Linux-5.15.167.4-microsoft-standard-WSL2-x86_64-with-glibc2.39)
     method: POST
     uri: https://management.azure.com/subscriptions/00000000-0000-0000-0000-000000000000/providers/Microsoft.SerialConsole/consoleServices/default/disableConsole?api-version=2024-07-01
   response:
@@ -29,7 +29,7 @@ interactions:
       content-type:
       - application/json; charset=UTF-8
       date:
-      - Thu, 21 Nov 2024 16:51:44 GMT
+      - Wed, 17 Dec 2025 12:49:02 GMT
       expires:
       - '-1'
       pragma:
@@ -42,12 +42,14 @@ interactions:
       - nosniff
       x-frame-options:
       - deny
+      x-ms-operation-identifier:
+      - tenantId=72f988bf-86f1-41af-91ab-2d7cd011db47,objectId=03039592-f720-4b7e-a2c4-fc19397bcba2/westeurope/0cf09d0b-d849-4c75-9bed-18a14590aa6f
       x-ms-ratelimit-remaining-subscription-global-writes:
       - '11999'
       x-ms-ratelimit-remaining-subscription-writes:
       - '799'
       x-msedge-ref:
-      - 'Ref A: 467FB350BAF9492C96DFB004512EE7F1 Ref B: MNZ221060609021 Ref C: 2024-11-21T16:51:44Z'
+      - 'Ref A: CFFAEE1F199D45EA98E8C1094AF0DE3A Ref B: AMS231032609017 Ref C: 2025-12-17T12:49:02Z'
     status:
       code: 200
       message: OK
@@ -67,7 +69,7 @@ interactions:
       Content-Type:
       - application/json
       User-Agent:
-      - AZURECLI/2.67.0 azsdk-python-core/1.31.0 Python/3.12.7 (Windows-11-10.0.26100-SP0)
+      - AZURECLI/2.81.0 azsdk-python-core/1.35.0 Python/3.12.3 (Linux-5.15.167.4-microsoft-standard-WSL2-x86_64-with-glibc2.39)
     method: POST
     uri: https://management.azure.com/subscriptions/00000000-0000-0000-0000-000000000000/providers/Microsoft.SerialConsole/consoleServices/default/enableConsole?api-version=2024-07-01
   response:
@@ -81,7 +83,7 @@ interactions:
       content-type:
       - application/json; charset=UTF-8
       date:
-      - Thu, 21 Nov 2024 16:51:45 GMT
+      - Wed, 17 Dec 2025 12:49:02 GMT
       expires:
       - '-1'
       pragma:
@@ -94,12 +96,14 @@ interactions:
       - nosniff
       x-frame-options:
       - deny
+      x-ms-operation-identifier:
+      - tenantId=72f988bf-86f1-41af-91ab-2d7cd011db47,objectId=03039592-f720-4b7e-a2c4-fc19397bcba2/westeurope/ff9774cd-b1cf-492a-82a2-a8c13852eb61
       x-ms-ratelimit-remaining-subscription-global-writes:
       - '11999'
       x-ms-ratelimit-remaining-subscription-writes:
       - '799'
       x-msedge-ref:
-      - 'Ref A: 01DB380606B146F6B7AA727925F4916A Ref B: MNZ221060618039 Ref C: 2024-11-21T16:51:45Z'
+      - 'Ref A: C96BCE73671C4907B60C62F22C76EB1B Ref B: AMS231022012021 Ref C: 2025-12-17T12:49:02Z'
     status:
       code: 200
       message: OK

src/serial-console/azext_serialconsole/tests/latest/recordings/test_check_resource_VMSS.yaml

@@ -197,7 +197,8 @@ def test_enable_disable(self):
 
 class SerialconsoleAdminCommandsTest(LiveScenarioTest):
 
-    def check_result(self, resource_group_name, vm_vmss_name, vmss_instanceid=None, message=""):
+    def check_result(self, resource_group_name, vm_vmss_name, vmss_instanceid=None, message="", hasManagedStorageAccount=False):
+        print("Checking serial console output for message: ", message)
         ARM_ENDPOINT = "https://management.azure.com"
         RP_PROVIDER = "Microsoft.SerialConsole"
         subscription_id = self.get_subscription_id()
@@ -215,21 +216,66 @@ def check_result(self, resource_group_name, vm_vmss_name, vmss_instanceid=None,
         headers = {'authorization': "Bearer " + access_token,
                    'accept': application_json_format,
                    'content-type': application_json_format}
-        result = requests.post(connection_url, headers=headers)
-        json_results = json.loads(result.text)
-        self.assertTrue(result.status_code ==
-                        200 and "connectionString" in json_results)
+
+        postRetryCounter = 1
+
+        while True:
+            result = requests.post(connection_url, headers=headers)
+            json_results = json.loads(result.text)
+
+            if result.status_code == 200 and "connectionString" in json_results:
+                break
+            else:
+                print("Failed to get connection string from serial console connect endpoint.")
+                print("Status code: ", result.status_code)
+                print("Response text: ", result.text)
+
+            if postRetryCounter > 3:
+                self.fail("Failed to get connection string from serial console connect endpoint after retrying multiple times... Failing test. See status and response of retries in logs.")
+            else:
+                postRetryCounter += 1
+                time.sleep(10)
+
         websocket_url = json_results["connectionString"]
 
         ws = websocket.WebSocket()
         ws.connect(websocket_url + "?authorization=" + access_token, timeout=30)
+        print("WebSocket connected for verifying message.")
+        print("Sleeping 60 seconds to allow 'Connecting...' message to appear.")
+        time.sleep(60)
+        
+        if not hasManagedStorageAccount:
+            print("Sending access token to start custom storage account setup...")
+            ws.send(access_token)
+            print("Finished sending access token")
+        
         buffer = ""
+        iter = 0
+        print("Starting to read from WebSocket to find message...")
         while True:
-            try:
-                buffer += ws.recv()
-            except (websocket.WebSocketTimeoutException, websocket.WebSocketConnectionClosedException):
+            iter += 1
+            print(f"Current timestamp: {time.strftime('%X')}, current iteration: {iter}/5")
+            
+            while True:
+                try:
+                    buffer += ws.recv()
+                except (websocket.WebSocketTimeoutException, websocket.WebSocketConnectionClosedException):
+                    break
+                        
+            if message in buffer:
+                print("Found message in buffer! Finished verification.")
                 break
-
+            
+            print(f"Message not found yet in buffer, current buffer: {buffer}")
+            
+            if iter >= 10:
+                print("Max retries reached, exiting read loop.")
+                break
+            else:
+                print("Sleeping 10 seconds before retrying...")
+                time.sleep(10) 
+        
+        ws.close()
         assert message in buffer
 
     @ResourceGroupPreparer(name_prefix='cli_test_serialconsole', location='westus2')
@@ -243,27 +289,13 @@ def test_send_sysrq_VMSS(self, resource_group, storage_account):
             'urn': 'Ubuntu2204',
             'loc': 'westus2'
         })
-        self.cmd(
-            'az vmss create -g {rg} -n {name} --image {urn} --instance-count 2 -l {loc} --orchestration-mode uniform')
-        self.cmd('az vmss update --name {name} --resource-group {rg} --set virtualMachineProfile.diagnosticsProfile="{{\\"bootDiagnostics\\": {{\\"Enabled\\" : \\"True\\",\\"StorageUri\\":\\"https://{sa}.blob.core.windows.net/\\"}}}}"')
-        result = self.cmd(
-            'vmss list-instances --resource-group {rg} --name {name} --query "[].instanceId"').get_output_in_json()
-        self.kwargs.update({'id': result[1]})
-        self.cmd(
-            'az vmss update-instances -g {rg} -n {name} --instance-ids {id}')
-        time.sleep(60)
-        for i in range(5):
-            try:
-                self.cmd('vmss get-instance-view --resource-group {rg} --name {name} --instance-id {id}', checks=[
-                    self.check('statuses[0].code',
-                               'ProvisioningState/succeeded'),
-                    self.check('statuses[1].code', 'PowerState/running'),
-                ])
-                break
-            except JMESPathCheckAssertionError:
-                time.sleep(30)
+        result = self.createVMSS()
+        print("Sending SysRq...")
+        stopwatch_start = time.time()
         self.cmd(
             'serial-console send sysrq -g {rg} -n {name} --instance-id {id} --input h')
+        stopwatch_end = time.time()
+        print(f"SysRq sent in {stopwatch_end - stopwatch_start} seconds.")
         self.check_result(resource_group, name,
                           vmss_instanceid=result[1], message="sysrq: HELP")
 
@@ -278,27 +310,13 @@ def test_send_nmi_VMSS(self, resource_group, storage_account):
             'urn': 'Ubuntu2204',
             'loc': 'westus2'
         })
-        self.cmd(
-            'az vmss create -g {rg} -n {name} --image {urn} --instance-count 2 -l {loc} --orchestration-mode uniform')
-        self.cmd('az vmss update --name {name} --resource-group {rg} --set virtualMachineProfile.diagnosticsProfile="{{\\"bootDiagnostics\\": {{\\"Enabled\\" : \\"True\\",\\"StorageUri\\":\\"https://{sa}.blob.core.windows.net/\\"}}}}"')
-        result = self.cmd(
-            'vmss list-instances --resource-group {rg} --name {name} --query "[].instanceId"').get_output_in_json()
-        self.kwargs.update({'id': result[1]})
-        self.cmd(
-            'az vmss update-instances -g {rg} -n {name} --instance-ids {id}')
-        time.sleep(60)
-        for i in range(5):
-            try:
-                self.cmd('vmss get-instance-view --resource-group {rg} --name {name} --instance-id {id}', checks=[
-                    self.check('statuses[0].code',
-                               'ProvisioningState/succeeded'),
-                    self.check('statuses[1].code', 'PowerState/running'),
-                ])
-                break
-            except JMESPathCheckAssertionError:
-                time.sleep(30)
+        result = self.createVMSS()
+        print("Sending NMI...")
+        stopwatch_start = time.time()
         self.cmd(
             'serial-console send nmi -g {rg} -n {name} --instance-id {id}')
+        stopwatch_end = time.time()
+        print(f"NMI sent in {stopwatch_end - stopwatch_start} seconds.")
         self.check_result(resource_group, name,
                           vmss_instanceid=result[1], message="NMI received")
 
@@ -313,26 +331,14 @@ def test_send_reset_VMSS(self, resource_group, storage_account):
             'urn': 'Ubuntu2204',
             'loc': 'westus2'
         })
-        self.cmd(
-            'az vmss create -g {rg} -n {name} --image {urn} --instance-count 2 -l {loc} --orchestration-mode uniform')
-        self.cmd('az vmss update --name {name} --resource-group {rg} --set virtualMachineProfile.diagnosticsProfile="{{\\"bootDiagnostics\\": {{\\"Enabled\\" : \\"True\\",\\"StorageUri\\":\\"https://{sa}.blob.core.windows.net/\\"}}}}"')
-        result = self.cmd(
-            'vmss list-instances --resource-group {rg} --name {name} --query "[].instanceId"').get_output_in_json()
-        self.kwargs.update({'id': result[1]})
-        self.cmd(
-            'az vmss update-instances -g {rg} -n {name} --instance-ids {id}')
-        time.sleep(60)
-        for i in range(5):
-            try:
-                self.cmd('vmss get-instance-view --resource-group {rg} --name {name} --instance-id {id}', checks=[
-                    self.check('statuses[0].code',
-                               'ProvisioningState/succeeded'),
-                    self.check('statuses[1].code', 'PowerState/running'),
-                ])
-                break
-            except JMESPathCheckAssertionError:
-                time.sleep(30)
+        result = self.createVMSS()
+        print("Sending Reset...")
+        stopwatch_start = time.time()
         self.cmd('serial-console send reset -g {rg} -n {name} --instance-id {id}')
+        stopwatch_end = time.time()
+        print(f"Reset sent in {stopwatch_end - stopwatch_start} seconds.")
+        self.check_result(resource_group, name,
+                          vmss_instanceid=result[1], message="Record successful boot")
 
     @ResourceGroupPreparer(name_prefix='cli_test_serialconsole', location='westus2')
     @StorageAccountPreparer(name_prefix='cli', location="westus2")
@@ -345,21 +351,12 @@ def test_send_nmi_VM(self, resource_group, storage_account):
             'urn': 'Ubuntu2204',
             'loc': 'westus2'
         })
-        self.cmd(
-            'az vm create -g {rg} -n {name} --image {urn} --boot-diagnostics-storage {sa} -l {loc} --generate-ssh-keys')
-        time.sleep(60)
-        for i in range(5):
-            try:
-                self.cmd('vm get-instance-view --resource-group {rg} --name {name}', checks=[
-                    self.check(
-                        'instanceView.statuses[0].code', 'ProvisioningState/succeeded'),
-                    self.check(
-                        'instanceView.statuses[1].code', 'PowerState/running'),
-                ])
-                break
-            except JMESPathCheckAssertionError:
-                time.sleep(30)
+        self.createVM()
+        print("Sending NMI...")
+        stopwatch_start = time.time()
         self.cmd('serial-console send nmi -g {rg} -n {name}')
+        stopwatch_end = time.time()
+        print(f"NMI sent in {stopwatch_end - stopwatch_start} seconds.")
         self.check_result(resource_group, name, message="NMI received")
 
     @ResourceGroupPreparer(name_prefix='cli_test_serialconsole', location='westus2')
@@ -373,21 +370,12 @@ def test_send_sysrq_VM(self, resource_group, storage_account):
             'urn': 'Ubuntu2204',
             'loc': 'westus2'
         })
-        self.cmd(
-            'az vm create -g {rg} -n {name} --image {urn} --boot-diagnostics-storage {sa} -l {loc} --generate-ssh-keys')
-        time.sleep(60)
-        for i in range(5):
-            try:
-                self.cmd('vm get-instance-view --resource-group {rg} --name {name}', checks=[
-                    self.check(
-                        'instanceView.statuses[0].code', 'ProvisioningState/succeeded'),
-                    self.check(
-                        'instanceView.statuses[1].code', 'PowerState/running'),
-                ])
-                break
-            except JMESPathCheckAssertionError:
-                time.sleep(30)
+        self.createVM()
+        print("Sending Sysrq...")
+        stopwatch_start = time.time()
         self.cmd('serial-console send sysrq -g {rg} -n {name} --input h')
+        stopwatch_end = time.time()
+        print(f"Sysrq sent in {stopwatch_end - stopwatch_start} seconds.")
         self.check_result(resource_group, name, message="sysrq: HELP")
 
     @ResourceGroupPreparer(name_prefix='cli_test_serialconsole', location='westus2')
@@ -401,6 +389,38 @@ def test_send_reset_VM(self, resource_group, storage_account):
             'urn': 'Ubuntu2204',
             'loc': 'westus2'
         })
+        self.createVM()
+        print("Sending Reset...")
+        stopwatch_start = time.time()
+        self.cmd('serial-console send reset -g {rg} -n {name}')
+        stopwatch_end = time.time()
+        print(f"Reset sent in {stopwatch_end - stopwatch_start} seconds.")
+        self.check_result(resource_group, name, message="Record successful boot")
+        
+    def createVMSS(self):
+        self.cmd(
+            'az vmss create -g {rg} -n {name} --image {urn} --instance-count 2 -l {loc} --orchestration-mode uniform')
+        self.cmd(
+            'az vmss update --name {name} --resource-group {rg} --set virtualMachineProfile.diagnosticsProfile="{{\\"bootDiagnostics\\": {{\\"Enabled\\" : \\"True\\",\\"StorageUri\\":\\"https://{sa}.blob.core.windows.net/\\"}}}}"')
+        result = self.cmd(
+            'vmss list-instances --resource-group {rg} --name {name} --query "[].instanceId"').get_output_in_json()
+        self.kwargs.update({'id': result[1]})
+        self.cmd(
+            'az vmss update-instances -g {rg} -n {name} --instance-ids {id}')
+        time.sleep(60)
+        for i in range(5):
+            try:
+                self.cmd('vmss get-instance-view --resource-group {rg} --name {name} --instance-id {id}', checks=[
+                    self.check('statuses[0].code',
+                               'ProvisioningState/succeeded'),
+                    self.check('statuses[1].code', 'PowerState/running'),
+                ])
+                break
+            except JMESPathCheckAssertionError:
+                time.sleep(30)
+        return result
+
+    def createVM(self):
         self.cmd(
             'az vm create -g {rg} -n {name} --image {urn} --boot-diagnostics-storage {sa} -l {loc} --generate-ssh-keys')
         time.sleep(60)
@@ -415,4 +435,3 @@ def test_send_reset_VM(self, resource_group, storage_account):
                 break
             except JMESPathCheckAssertionError:
                 time.sleep(30)
-        self.cmd('serial-console send reset -g {rg} -n {name}')

src/serial-console/azext_serialconsole/tests/latest/recordings/test_enable_disable.yaml

@@ -16,7 +16,7 @@
 
 # TODO: Confirm this is the right version number you want and it matches your
 # HISTORY.rst entry.
-VERSION = '1.0.0b2'
+VERSION = '1.0.0b3'
 
 # The full list of classifiers is available at
 # https://pypi.python.org/pypi?%3Aaction=list_classifiers