Added support for track 2 SDKs (azure-mgmt-iothub >= 1.0.0) (#329)

azext_iot/common/utility.py

@@ -438,10 +438,14 @@ def is_iso8601_time(self, to_validate: str) -> bool:
             return False
 
 
-def ensure_min_version(cur_ver, min_ver):
-    from pkg_resources._vendor.packaging import version
+def ensure_iothub_sdk_min_version(min_ver):
+    from packaging import version
+    try:
+        from azure.mgmt.iothub import __version__ as iot_sdk_version
+    except ImportError:
+        from azure.mgmt.iothub._configuration import VERSION as iot_sdk_version
 
-    return version.parse(cur_ver) >= version.parse(min_ver)
+    return version.parse(iot_sdk_version) >= version.parse(min_ver)
 
 
 def scantree(path):

azext_iot/constants.py

@@ -48,3 +48,6 @@
 
 # Config Key's
 CONFIG_KEY_UAMQP_EXT_VERSION = "uamqp_ext_version"
+
+# Initial Track 2 SDK version
+IOTHUB_TRACK_2_SDK_MIN_VERSION = '1.0.0'

azext_iot/iothub/providers/base.py

@@ -21,7 +21,7 @@ def __init__(self, cmd, hub_name, rg, login=None):
         self.discovery = IotHubDiscovery(cmd)
         self.target = self.discovery.get_target(
             hub_name=self.hub_name,
-            rg=self.rg,
+            resource_group_name=self.rg,
             login=login,
         )
         self.resolver = SdkResolver(self.target)

azext_iot/iothub/providers/discovery.py

@@ -6,10 +6,13 @@
 
 from knack.util import CLIError
 from knack.log import get_logger
-from azext_iot.common.utility import trim_from_start
+from azure.cli.core.commands.client_factory import get_subscription_id
+from azext_iot.common.utility import trim_from_start, ensure_iothub_sdk_min_version
 from azext_iot.iothub.models.iothub_target import IotHubTarget
 from azext_iot._factory import iot_hub_service_factory
+from azext_iot.constants import IOTHUB_TRACK_2_SDK_MIN_VERSION
 from typing import Dict, List
+from enum import Enum, EnumMeta
 
 PRIVILEDGED_ACCESS_RIGHTS_SET = set(
     ["RegistryWrite", "ServiceConnect", "DeviceConnect"]
@@ -30,9 +33,9 @@ def _initialize_client(self):
         if not self.client:
             if getattr(self.cmd, "cli_ctx", None):
                 self.client = iot_hub_service_factory(self.cmd.cli_ctx)
+                self.sub_id = get_subscription_id(self.cmd.cli_ctx)
             else:
                 self.client = self.cmd
-            self.sub_id = self.client.config.subscription_id
 
     def get_iothubs(self, rg: str = None) -> List:
         self._initialize_client()
@@ -44,11 +47,15 @@ def get_iothubs(self, rg: str = None) -> List:
         else:
             hubs_pager = self.client.list_by_resource_group(resource_group_name=rg)
 
-        try:
-            while True:
-                hubs_list.extend(hubs_pager.advance_page())
-        except StopIteration:
-            pass
+        if ensure_iothub_sdk_min_version(IOTHUB_TRACK_2_SDK_MIN_VERSION):
+            for hubs in hubs_pager.by_page():
+                hubs_list.extend(hubs)
+        else:
+            try:
+                while True:
+                    hubs_list.extend(hubs_pager.advance_page())
+            except StopIteration:
+                pass
 
         return hubs_list
 
@@ -60,23 +67,25 @@ def get_policies(self, hub_name: str, rg: str) -> List:
         )
         policy_list = []
 
-        try:
-            while True:
-                policy_list.extend(policy_pager.advance_page())
-        except StopIteration:
-            pass
+        if ensure_iothub_sdk_min_version(IOTHUB_TRACK_2_SDK_MIN_VERSION):
+            for policy in policy_pager.by_page():
+                policy_list.extend(policy)
+        else:
+            try:
+                while True:
+                    policy_list.extend(policy_pager.advance_page())
+            except StopIteration:
+                pass
 
         return policy_list
 
     def find_iothub(self, hub_name: str, rg: str = None):
         self._initialize_client()
 
-        from azure.mgmt.iothub.models import ErrorDetailsException
-
         if rg:
             try:
                 return self.client.get(resource_group_name=rg, resource_name=hub_name)
-            except ErrorDetailsException:
+            except:  # pylint: disable=broad-except
                 raise CLIError(
                     "Unable to find IoT Hub: {} in resource group: {}".format(
                         hub_name, rg
@@ -128,12 +137,12 @@ def find_policy(self, hub_name: str, rg: str, policy_name: str = "auto"):
     def get_target_by_cstring(cls, connection_string: str) -> IotHubTarget:
         return IotHubTarget.from_connection_string(cstring=connection_string).as_dict()
 
-    def get_target(self, hub_name: str, rg: str = None, **kwargs) -> Dict[str, str]:
+    def get_target(self, hub_name: str, resource_group_name: str = None, **kwargs) -> Dict[str, str]:
         cstring = kwargs.get("login")
         if cstring:
             return self.get_target_by_cstring(connection_string=cstring)
 
-        target_iothub = self.find_iothub(hub_name=hub_name, rg=rg)
+        target_iothub = self.find_iothub(hub_name=hub_name, rg=resource_group_name)
 
         policy_name = kwargs.get("policy_name", "auto")
         rg = target_iothub.additional_properties.get("resourcegroup")
@@ -151,13 +160,13 @@ def get_target(self, hub_name: str, rg: str = None, **kwargs) -> Dict[str, str]:
             include_events=include_events,
         )
 
-    def get_targets(self, rg: str = None, **kwargs) -> List[Dict[str, str]]:
+    def get_targets(self, resource_group_name: str = None, **kwargs) -> List[Dict[str, str]]:
         targets = []
-        hubs = self.get_iothubs(rg=rg)
+        hubs = self.get_iothubs(rg=resource_group_name)
         if hubs:
             for hub in hubs:
                 targets.append(
-                    self.get_target(hub_name=hub.name, rg=self._get_rg(hub), **kwargs)
+                    self.get_target(hub_name=hub.name, resource_group_name=self._get_rg(hub), **kwargs)
                 )
 
         return targets
@@ -186,7 +195,7 @@ def _build_target(
         target["subscription"] = self.sub_id
         target["resourcegroup"] = iothub.additional_properties.get("resourcegroup")
         target["location"] = iothub.location
-        target["sku_tier"] = iothub.sku.tier.value
+        target["sku_tier"] = iothub.sku.tier.value if isinstance(iothub.sku.tier, (Enum, EnumMeta)) else iothub.sku.tier
 
         if include_events:
             events = {}

azext_iot/operations/hub.py

@@ -9,6 +9,7 @@
 import six
 from knack.log import get_logger
 from knack.util import CLIError
+from enum import Enum, EnumMeta
 from azext_iot.constants import (
     EXTENSION_ROOT,
     DEVICE_DEVICESCOPE_PREFIX,
@@ -36,7 +37,7 @@
     unpack_msrest_error,
     init_monitoring,
     process_json_arg,
-    ensure_min_version,
+    ensure_iothub_sdk_min_version,
     generate_key
 )
 from azext_iot._factory import SdkResolver, CloudError
@@ -55,7 +56,7 @@ def iot_query(
     top = _process_top(top)
     discovery = IotHubDiscovery(cmd)
     target = discovery.get_target(
-        hub_name=hub_name, rg=resource_group_name, login=login
+        hub_name=hub_name, resource_group_name=resource_group_name, login=login
     )
     resolver = SdkResolver(target=target)
     service_sdk = resolver.get_sdk(SdkType.service_sdk)
@@ -2287,7 +2288,6 @@ def iot_device_export(
     resource_group_name=None,
 ):
     from azext_iot._factory import iot_hub_service_factory
-    from azure.mgmt.iothub import __version__ as iot_sdk_version
 
     client = iot_hub_service_factory(cmd.cli_ctx)
     discovery = IotHubDiscovery(cmd)
@@ -2298,7 +2298,7 @@ def iot_device_export(
     if exists(blob_container_uri):
         blob_container_uri = read_file_content(blob_container_uri)
 
-    if ensure_min_version(iot_sdk_version, "0.12.0"):
+    if ensure_iothub_sdk_min_version("0.12.0"):
         from azure.mgmt.iothub.models import ExportDevicesRequest
         from azext_iot.common.shared import AuthenticationType
 
@@ -2336,7 +2336,6 @@ def iot_device_import(
     resource_group_name=None,
 ):
     from azext_iot._factory import iot_hub_service_factory
-    from azure.mgmt.iothub import __version__ as iot_sdk_version
 
     client = iot_hub_service_factory(cmd.cli_ctx)
     discovery = IotHubDiscovery(cmd)
@@ -2350,7 +2349,7 @@ def iot_device_import(
     if exists(output_blob_container_uri):
         output_blob_container_uri = read_file_content(output_blob_container_uri)
 
-    if ensure_min_version(iot_sdk_version, "0.12.0"):
+    if ensure_iothub_sdk_min_version("0.12.0"):
         from azure.mgmt.iothub.models import ImportDevicesRequest
         from azext_iot.common.shared import AuthenticationType
 
@@ -2711,7 +2710,7 @@ def _get_hub_connection_string(
                 entityPath,
             )
             for p in policies
-            if "serviceconnect" in p.rights.value.lower()
+            if "serviceconnect" in (p.rights.value.lower() if isinstance(p.rights, (Enum, EnumMeta)) else p.rights.lower())
         ]
 
     hostname = hub.properties.host_name

azext_iot/tests/iothub/jobs/test_iothub_jobs_int.py

@@ -182,6 +182,21 @@ def test_jobs(self):
             checks=[self.check("jobId", self.job_ids[2])],
         )
 
+        # Allow time for job to transfer to scheduled state (cannot cancel job in running state)
+        from time import sleep
+        sleep(5)
+
+        self.cmd(
+            "iot hub job show --job-id {} -n {} -g {}".format(
+                self.job_ids[2], LIVE_HUB, LIVE_RG
+            ),
+            checks=[
+                self.check("jobId", self.job_ids[2]),
+                self.check("status", "scheduled"),
+            ],
+        )
+
+        # Cancel job
         self.cmd(
             "iot hub job cancel --job-id {} -n {} -g {}".format(
                 self.job_ids[2], LIVE_HUB, LIVE_RG

azext_iot/tests/iothub/test_iothub_discovery_int.py

@@ -74,7 +74,7 @@ def test_iothub_targets(self):
         auto_target = discovery.get_target(hub_name=LIVE_HUB)
         assert_target(auto_target, rg=LIVE_RG)
 
-        auto_target = discovery.get_target(hub_name=LIVE_HUB, rg=LIVE_RG)
+        auto_target = discovery.get_target(hub_name=LIVE_HUB, resource_group_name=LIVE_RG)
         assert_target(auto_target, rg=LIVE_RG)
 
         desired_target = discovery.get_target(
@@ -85,7 +85,7 @@ def test_iothub_targets(self):
         sub_targets = discovery.get_targets()
         [assert_target(tar) for tar in sub_targets]
 
-        rg_targets = discovery.get_targets(rg=LIVE_RG, include_events=True)
+        rg_targets = discovery.get_targets(resource_group_name=LIVE_RG, include_events=True)
         [assert_target(tar, rg=LIVE_RG, include_events=True) for tar in rg_targets]
 
         assert len(rg_targets) <= len(sub_targets)
@@ -99,7 +99,7 @@ def assert_target(target: dict, by_cstring=False, include_events=False, **kwargs
 
     if not by_cstring:
         assert target["secondarykey"]
-        assert target["subscription"]
+        assert target["subscription"] and target["subscription"] != "unknown"
 
         if "rg" in kwargs:
             assert target["resourcegroup"] == kwargs["rg"]

azext_iot/tests/utility/test_iot_utility_unit.py

@@ -16,7 +16,7 @@
     process_json_arg,
     read_file_content,
     logger,
-    ensure_min_version,
+    ensure_iothub_sdk_min_version,
 )
 from azext_iot.common.deps import ensure_uamqp
 from azext_iot.constants import EVENT_LIB, EXTENSION_NAME
@@ -294,8 +294,13 @@ class TestVersionComparison(object):
             ("2.0.1.9", "2.0.6", False),
         ],
     )
-    def test_ensure_min_version(self, current, minimum, expected):
-        assert ensure_min_version(current, minimum) == expected
+    def test_ensure_iothub_sdk_min_version(self, mocker, current, minimum, expected):
+        try:
+            mocker.patch("azure.mgmt.iothub.__version__", current)
+        except:
+            mocker.patch("azure.mgmt.iothub._configuration.VERSION", current)
+
+        assert ensure_iothub_sdk_min_version(minimum) == expected
 
 
 class TestEmbeddedCli(object):

setup.py

@@ -41,14 +41,14 @@
 # 'jmespath==0.9.3',
 # 'pyyaml==3.13'
 # 'knack>=0.3.1'
-# 'jsonschema==3.0.2'
+# 'jsonschema==3.2.0'
 # 'enum34' (when python_version < 3.4)
 
 # There is also a dependency for uamqp for amqp based commands
 # though that is installed out of band (managed by the extension)
 # for compatibility reasons.
 
-DEPENDENCIES = ["paho-mqtt==1.5.0", "jsonschema==3.2.0", "setuptools"]
+DEPENDENCIES = ["paho-mqtt==1.5.0", "jsonschema==3.2.0", "packaging"]
 
 
 CLASSIFIERS = [