[Bug Fix] Fix "cannot unpack non-iterable NoneType object" error in OS check

src/ssh/HISTORY.md

@@ -1,10 +1,13 @@
 Release History
 ===============
+2.0.2
+-----
+* [Bug Fix] Fix logic that checks for the OS of the target machine to avoid "cannot unpack non-iterable NoneType object" error
+
 2.0.1
 -----
 * [Bug fix] For connections to arc resources, stop attempting to create new service configuration if user has no permission to read service configuration.
 
-
 2.0.0
 -----
 * [BREAKING CHANGE] Update Microsoft.HybridConnectivity SDK to stable version, which adds functionality to enable SSH connections on specified ports in your Arc Server using an API, instead of enabling ports for connection locally in the Arc Agent running in the target machine. New connections might fail after updating the extension, since the port for connection will need to be enabled at the HybridConnectivity Resource Provider at the first connection attempt. This change doesn't affect those who use this extension to connect to Azure Virtual Machines.

src/ssh/azext_ssh/target_os_utils.py

@@ -80,10 +80,12 @@ def _get_arc_server_os(cmd, resource_group_name, vm_name):
     try:
         arc_server = ArcServerShow(cli_ctx=cmd.cli_ctx)(command_args=get_args)
     except Exception:
-        return None
+        return None, None
 
     if arc_server and arc_server.get('osName', None):
-        os_type = arc_server['osName']
+        os_type = arc_server.get('osName')
+    elif arc_server and arc_server.get('properties', None):
+        os_type = arc_server.get('properties').get('osType', None)
 
     if arc_server and arc_server.get('properties'):
         agent_version = arc_server.get('properties').get('agentVersion')
@@ -104,7 +106,7 @@ def _get_connected_vmware_os(cmd, resource_group_name, vm_name):
     try:
         vmware = VMwarevSphereShow(cli_ctx=cmd.cli_ctx)(command_args=get_args)
     except Exception:
-        return None
+        return None, None
 
     if vmware and vmware.get("osProfile") and vmware.get("osProfile").get("osType"):
         os_type = vmware.get("osProfile").get("osType")

src/ssh/azext_ssh/tests/latest/test_target_os_utils.py

@@ -0,0 +1,91 @@
+# --------------------------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for license information.
+# --------------------------------------------------------------------------------------------
+
+import unittest
+from unittest import mock
+
+from azext_ssh import target_os_utils
+
+class TargetOSUtilsTest(unittest.TestCase):
+
+    @mock.patch('azext_ssh.aaz.latest.hybrid_compute.machine.Show')
+    def test_get_arc_os(self, mock_get_arc):
+        cmd = mock.Mock()
+        cmd.cli_ctx = mock.Mock()
+
+        showclass = mock.Mock()
+        showclass.return_value = {
+            "properties": {
+                "osType": "os_type",
+                "agentVersion": "arc_agent_version"
+            }
+        }
+
+        mock_get_arc.return_value = showclass
+
+        os, agent = target_os_utils._get_arc_server_os(cmd, "rg", "vm")
+
+        self.assertEqual(os, "os_type")
+        self.assertEqual(agent, "arc_agent_version")
+
+
+
+    @mock.patch('azext_ssh.aaz.latest.hybrid_compute.machine.Show', autospec=True)
+    def test_get_arc_os_exception(self, mock_get_arc):
+        cmd = mock.Mock()
+        cmd.cli_ctx = mock.Mock()
+
+        mock_get_arc.return_value.side_effect = mock.Mock(side_effect=Exception('Test'))
+
+        os, agent = target_os_utils._get_arc_server_os(cmd, "rg", "vm")
+
+        self.assertEqual(os, None)
+        self.assertEqual(agent, None)
+
+
+    @mock.patch('azext_ssh.aaz.latest.connected_v_mwarev_sphere.virtual_machine.Show')
+    def test_get_vmware_os(self, mock_get_vmware):
+        cmd = mock.Mock()
+        cmd.cli_ctx = mock.Mock()
+
+        showclass = mock.Mock()
+        showclass.return_value = {
+            "osProfile":{
+                "osType": "os_type"
+            },
+            "properties":{
+                "guestAgentProfile": {
+                    "agentVersion": "agent_version"
+                }
+            }
+        }
+
+        mock_get_vmware.return_value = showclass
+
+        os, agent = target_os_utils._get_connected_vmware_os(cmd, "rg", "vm")
+
+        self.assertEqual(os, "os_type")
+        self.assertEqual(agent, "agent_version")
+
+
+
+    @mock.patch('azext_ssh.aaz.latest.connected_v_mwarev_sphere.virtual_machine.Show', autospec=True)
+    def test_get_vmware_os_exception(self, mock_get_vmware):
+        cmd = mock.Mock()
+        cmd.cli_ctx = mock.Mock()
+
+        mock_get_vmware.return_value.side_effect = mock.Mock(side_effect=Exception('Test'))
+
+        os, agent = target_os_utils._get_connected_vmware_os(cmd, "rg", "vm")
+
+        self.assertEqual(os, None)
+        self.assertEqual(agent, None)
+
+
+
+
+
+if __name__ == '__main__':
+    unittest.main()

src/ssh/setup.py

@@ -7,7 +7,7 @@
 
 from setuptools import setup, find_packages
 
-VERSION = "2.0.1"
+VERSION = "2.0.2"
 
 CLASSIFIERS = [
     'Development Status :: 4 - Beta',