Updated to support multihub in iot central

azext_iot/common/_azure.py

@@ -241,15 +241,21 @@ def _find_iot_dps_from_list(all_dps, dps_name):
 
 
 def get_iot_central_tokens(cmd, app_id, central_api_uri):
-    def get_event_hub_token(app_id, iotcAccessToken):
+    def formatUrl(url):
+        finalUrl = url
+        if not finalUrl.startswith('https://'):
+            finalUrl = 'https://{}'.format(finalUrl)
+        return finalUrl
+
+    def get_tokens(app_id, iotcAccessToken, central_api_uri):
         import requests
-        url = "https://{}/v1-beta/applications/{}/diagnostics/sasTokens".format(central_api_uri, app_id)
+        url = "{}.{}/system/iothubs/generateSasTokens".format(formatUrl(app_id), central_api_uri)
         response = requests.post(url, headers={'Authorization': 'Bearer {}'.format(iotcAccessToken)})
         return response.json()
 
     aad_token = _get_aad_token(cmd, resource="https://apps.azureiotcentral.com")['accessToken']
 
-    tokens = get_event_hub_token(app_id, aad_token)
+    tokens = get_tokens(app_id, aad_token, central_api_uri)
 
     if tokens.get('error'):
         raise CLIError(
@@ -259,10 +265,6 @@ def get_event_hub_token(app_id, iotcAccessToken):
     return tokens
 
 
-def get_iot_hub_token_from_central_app_id(cmd, app_id, central_api_uri):
-    return get_iot_central_tokens(cmd, app_id, central_api_uri)['iothubTenantSasToken']['sasToken']
-
-
 def get_iot_pnp_connection_string(
         cmd,
         endpoint,

azext_iot/operations/central.py

@@ -6,7 +6,6 @@
 
 from knack.util import CLIError
 from azext_iot._factory import _bind_sdk
-from azext_iot.common._azure import get_iot_hub_token_from_central_app_id
 from azext_iot.common.shared import SdkType
 from azext_iot.common.utility import unpack_msrest_error, init_monitoring
 from azext_iot.common.sas_token_auth import BasicSasTokenAuthentication
@@ -16,18 +15,26 @@ def find_between(s, start, end):
     return (s.split(start))[1].split(end)[0]
 
 
-def iot_central_device_show(
-    cmd, device_id, app_id, central_api_uri="api.azureiotcentral.com"
-):
-    sasToken = get_iot_hub_token_from_central_app_id(cmd, app_id, central_api_uri)
-    endpoint = find_between(sasToken, "SharedAccessSignature sr=", "&sig=")
-    target = {"entity": endpoint}
-    auth = BasicSasTokenAuthentication(sas_token=sasToken)
-    service_sdk, errors = _bind_sdk(target, SdkType.service_sdk, auth=auth)
-    try:
-        return service_sdk.get_twin(device_id)
-    except errors.CloudError as e:
-        raise CLIError(unpack_msrest_error(e))
+def iot_central_device_show(cmd, device_id, app_id, central_api_uri='azureiotcentral.com'):
+    from azext_iot.common._azure import get_iot_central_tokens
+
+    tokens = get_iot_central_tokens(cmd, app_id, central_api_uri)
+    exception = None
+
+    # try all hubs, return first exception if none succeed
+    for key in tokens:
+        sasToken = tokens[key]['iothubTenantSasToken']['sasToken']
+        endpoint = find_between(sasToken, 'SharedAccessSignature sr=', '&sig=')
+        target = {'entity': endpoint}
+        auth = BasicSasTokenAuthentication(sas_token=sasToken)
+        service_sdk, errors = _bind_sdk(target, SdkType.service_sdk, auth=auth)
+        try:
+            return service_sdk.get_twin(device_id)
+        except errors.CloudError as e:
+            if(exception is None):
+                exception = CLIError(unpack_msrest_error(e))
+
+    raise exception
 
 
 def iot_central_validate_messages(
@@ -41,7 +48,7 @@ def iot_central_validate_messages(
     repair=False,
     properties=None,
     yes=False,
-    central_api_uri="api.azureiotcentral.com",
+    central_api_uri="azureiotcentral.com",
 ):
     _events3_runner(
         cmd,
@@ -69,7 +76,7 @@ def iot_central_monitor_events(
     repair=False,
     properties=None,
     yes=False,
-    central_api_uri="api.azureiotcentral.com",
+    central_api_uri="azureiotcentral.com",
 ):
     _events3_runner(
         cmd,
@@ -107,18 +114,15 @@ def _events3_runner(
 
     from azext_iot.operations.events3 import _builders, _events
 
-    eventHubTarget = _builders.EventTargetBuilder().build_central_event_hub_target(
-        cmd, app_id, central_api_uri
-    )
+    eventHubTargets = _builders.EventTargetBuilder().build_central_event_hub_target(cmd, app_id, central_api_uri)
+    executorTargets = []
 
-    _events.executor(
-        eventHubTarget,
-        consumer_group=consumer_group,
-        enqueued_time=enqueued_time,
-        properties=properties,
-        timeout=timeout,
-        device_id=device_id,
-        output=output,
-        validate_messages=validate_messages,
-        simulate_errors=simulate_errors,
-    )
+    for target in eventHubTargets:
+        executorTargets.append(_events.executorData(target, consumer_group))
+
+    _events.nExecutor(executorTargets,
+                      enqueued_time=enqueued_time,
+                      properties=properties,
+                      timeout=timeout,
+                      device_id=device_id,
+                      output=output)

azext_iot/operations/events3/_builders.py

@@ -79,30 +79,36 @@ async def _evaluate_redirect(self, endpoint):
     async def _build_central_event_hub_target_async(self, cmd, app_id, central_api_uri):
         from azext_iot.common._azure import get_iot_central_tokens
 
-        tokens = get_iot_central_tokens(cmd, app_id, central_api_uri)
-        eventHubToken = tokens['eventhubSasToken']
-        hostnameWithoutPrefix = eventHubToken['hostname'].split("/")[2]
-        endpoint = hostnameWithoutPrefix
-        path = eventHubToken["entityPath"]
-        tokenExpiry = tokens['expiry']
-        auth = self._build_auth_container_from_token(endpoint, path, eventHubToken['sasToken'], tokenExpiry)
-        address = "amqps://{}/{}/$management".format(hostnameWithoutPrefix, eventHubToken["entityPath"])
-        meta_data = await self._query_meta_data(address, path, auth)
-        partition_count = meta_data[b'partition_count']
-        partition_ids = []
-        for i in range(int(partition_count)):
-            partition_ids.append(str(i))
-        partitions = partition_ids
-        auth = self._build_auth_container_from_token(endpoint, path, eventHubToken['sasToken'], tokenExpiry)
+        allTokens = get_iot_central_tokens(cmd, app_id, central_api_uri)
+        targets = []
+        # create target for each event hub
+        for key in allTokens:
+            tokens = allTokens[key]
+            eventHubToken = tokens['eventhubSasToken']
+            hostnameWithoutPrefix = eventHubToken['hostname'].split("/")[2]
+            endpoint = hostnameWithoutPrefix
+            path = eventHubToken["entityPath"]
+            tokenExpiry = tokens['expiry']
+            auth = self._build_auth_container_from_token(endpoint, path, eventHubToken['sasToken'], tokenExpiry)
+            address = "amqps://{}/{}/$management".format(hostnameWithoutPrefix, eventHubToken["entityPath"])
+            meta_data = await self._query_meta_data(address, path, auth)
+            partition_count = meta_data[b'partition_count']
+            partition_ids = []
+            for i in range(int(partition_count)):
+                partition_ids.append(str(i))
+            partitions = partition_ids
+            auth = self._build_auth_container_from_token(endpoint, path, eventHubToken['sasToken'], tokenExpiry)
 
-        eventHubTarget = {
-            'endpoint': endpoint,
-            'path': path,
-            'auth': auth,
-            'partitions': partitions
-        }
+            eventHubTarget = {
+                'endpoint': endpoint,
+                'path': path,
+                'auth': auth,
+                'partitions': partitions
+            }
 
-        return eventHubTarget
+            targets.append(eventHubTarget)
+
+        return targets
 
     async def _build_iot_hub_target_async(self, target):
         if 'events' not in target:

azext_iot/operations/events3/_events.py

@@ -24,6 +24,14 @@
 logger = get_logger(__name__)
 
 
+class executorData:
+    def __init__(self,
+                 target,
+                 consumer_group):
+        self.target = target
+        self.consumer_group = consumer_group
+
+
 def executor(
     target,
     consumer_group,
@@ -39,25 +47,54 @@ def executor(
     validate_messages=False,
     simulate_errors=False,
 ):
-
+    executor = executorData(target, consumer_group)
+
+    return nExecutor([executor], enqueued_time,
+                     properties,
+                     timeout,
+                     device_id,
+                     output,
+                     content_type,
+                     devices,
+                     interface_name,
+                     pnp_context,
+                     validate_messages,
+                     simulate_errors)
+
+
+def nExecutor(
+    executorTargets,
+    enqueued_time,
+    properties=None,
+    timeout=0,
+    device_id=None,
+    output=None,
+    content_type=None,
+    devices=None,
+    interface_name=None,
+    pnp_context=None,
+    validate_messages=False,
+    simulate_errors=False,
+):
     coroutines = []
-    coroutines.append(
-        initiate_event_monitor(
-            target,
-            consumer_group,
-            enqueued_time,
-            device_id,
-            properties,
-            timeout,
-            output,
-            content_type,
-            devices,
-            interface_name,
-            pnp_context,
-            validate_messages,
-            simulate_errors,
+    for executor in executorTargets:
+        coroutines.append(
+            initiate_event_monitor(
+                executor.target,
+                executor.consumer_group,
+                enqueued_time,
+                device_id,
+                properties,
+                timeout,
+                output,
+                content_type,
+                devices,
+                interface_name,
+                pnp_context,
+                validate_messages,
+                simulate_errors
+            )
         )
-    )
 
     loop = asyncio.get_event_loop()
     if loop.is_closed():

azext_iot/tests/test_iot_central_unit.py

@@ -20,17 +20,6 @@
 resource = "shared_resource"
 
 
-@pytest.fixture()
-def fixture_iot_token(mocker):
-    sas = mocker.patch(
-        "azext_iot.operations.central.get_iot_hub_token_from_central_app_id"
-    )
-    sas.return_value = "SharedAccessSignature sr={}&sig=signature&se=expiry&skn=service".format(
-        resource
-    )
-    return sas
-
-
 @pytest.fixture()
 def fixture_cmd(mocker):
     # Placeholder for later use
@@ -96,13 +85,15 @@ def fixture_get_iot_central_tokens(mocker):
     mock = mocker.patch("azext_iot.common._azure.get_iot_central_tokens")
 
     mock.return_value = {
-        "eventhubSasToken": {
-            "hostname": "part1/part2/part3",
-            "entityPath": "entityPath",
-            "sasToken": "sasToken",
-        },
-        "expiry": "0000",
-        "iothubTenantSasToken": {"sasToken": "iothubTenantSasToken"},
+        "id": {
+            "eventhubSasToken": {
+                "hostname": "part1/part2/part3",
+                "entityPath": "entityPath",
+                "sasToken": "sasToken",
+            },
+            "expiry": "0000",
+            "iothubTenantSasToken": {"sasToken": "SharedAccessSignature sr=shared_resource&sig="},
+        }
     }
 
 
@@ -131,24 +122,13 @@ class Cmd:
             "tokenType": "raw token 0 - A",
         }
 
-    def test_get_iot_hub_token_from_central_app_id(
-        self, fixture_get_iot_central_tokens
-    ):
-        from azext_iot.common._azure import get_iot_hub_token_from_central_app_id
-
-        # Test to ensure get_iot_hub_token_from_central_app_id returns iothubTenantSasToken
-        assert (
-            get_iot_hub_token_from_central_app_id({}, "app_id", "api-uri")
-            == "iothubTenantSasToken"
-        )
-
 
 class TestDeviceTwinShow:
     def test_device_twin_show_calls_get_twin(
-        self, fixture_iot_token, fixture_bind_sdk, fixture_cmd
+        self, fixture_bind_sdk, fixture_cmd, fixture_get_iot_central_tokens
     ):
         result = subject.iot_central_device_show(
-            fixture_cmd, device_id, app_id, "api-uri"
+            fixture_cmd, device_id, app_id
         )
 
         # Ensure get_twin is called and result is returned