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

Author: Razvan Ivan

src/serial-console/HISTORY.rst

@@ -1,5 +1,9 @@
 Release History
 ===============
+1.0.1
+++++++
+* 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/test_serialconsole_scenario.py

@@ -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()
@@ -217,19 +218,53 @@ def check_result(self, resource_group_name, vm_vmss_name, vmss_instanceid=None,
                    '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)
+        
+        if (not result.status_code == 200) or ("connectionString" not in json_results):
+            print("Failed to get connection string from serial console connect endpoint.")
+            print("Status code: ", result.status_code)
+            print("Response text: ", result.text)
+            self.fail("Failed to get connection string from serial console connect endpoint. See status and response in logs.")
+        
         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 +278,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 +299,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 +320,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 +340,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 +359,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 +378,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 +424,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/setup.py

@@ -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.1'
 
 # The full list of classifiers is available at
 # https://pypi.python.org/pypi?%3Aaction=list_classifiers