From 33f7cef82d8b0034a7a6b667e8dea8e5212b8ae6 Mon Sep 17 00:00:00 2001
From: Yunhao Ling <adam_ling@outlook.com>
Date: Mon, 19 Aug 2019 20:37:34 -0700
Subject: [PATCH 1/5] Fix bug that iothub hub can't receive

---
 .../azure/eventhub/aio/client_async.py               |  6 +++++-
 .../azure/eventhub/aio/consumer_async.py             | 11 ++++++++---
 .../azure-eventhubs/azure/eventhub/client.py         |  7 ++++++-
 .../azure-eventhubs/azure/eventhub/consumer.py       | 12 +++++++++---
 4 files changed, 28 insertions(+), 8 deletions(-)

diff --git a/sdk/eventhub/azure-eventhubs/azure/eventhub/aio/client_async.py b/sdk/eventhub/azure-eventhubs/azure/eventhub/aio/client_async.py
index 4ad5a99f538c..a6be575f43d8 100644
--- a/sdk/eventhub/azure-eventhubs/azure/eventhub/aio/client_async.py
+++ b/sdk/eventhub/azure-eventhubs/azure/eventhub/aio/client_async.py
@@ -101,10 +101,14 @@ async def _close_connection(self):
         await self._conn_manager.reset_connection_if_broken()
 
     async def _management_request(self, mgmt_msg, op_type):
+        alt_creds = {
+            "username": self._auth_config.get("iot_username"),
+            "password": self._auth_config.get("iot_password")
+        }
         max_retries = self.config.max_retries
         retry_count = 0
         while True:
-            mgmt_auth = self._create_auth()
+            mgmt_auth = self._create_auth(**alt_creds)
             mgmt_client = AMQPClientAsync(self.mgmt_target, auth=mgmt_auth, debug=self.config.network_tracing)
             try:
                 conn = await self._conn_manager.get_connection(self.host, mgmt_auth)
diff --git a/sdk/eventhub/azure-eventhubs/azure/eventhub/aio/consumer_async.py b/sdk/eventhub/azure-eventhubs/azure/eventhub/aio/consumer_async.py
index 404fc23312f0..6ca674320071 100644
--- a/sdk/eventhub/azure-eventhubs/azure/eventhub/aio/consumer_async.py
+++ b/sdk/eventhub/azure-eventhubs/azure/eventhub/aio/consumer_async.py
@@ -110,8 +110,9 @@ async def __anext__(self):
 
     def _create_handler(self):
         alt_creds = {
-            "username": self.client._auth_config.get("iot_username"),
-            "password": self.client._auth_config.get("iot_password")}
+            "username": self.client._auth_config.get("iot_username") if self.redirected else None,
+            "password": self.client._auth_config.get("iot_password") if self.redirected else None
+        }
         source = Source(self.source)
         if self.offset is not None:
             source.set_filter(self.offset._selector())
@@ -134,7 +135,7 @@ async def _redirect(self, redirect):
         self.messages_iter = None
         await super(EventHubConsumer, self)._redirect(redirect)
 
-    async def _open(self, timeout_time=None):
+    async def _open(self, timeout_time=None, **kwargs):
         """
         Open the EventHubConsumer using the supplied connection.
         If the handler has previously been redirected, the redirect
@@ -147,6 +148,10 @@ async def _open(self, timeout_time=None):
             self.source = self.redirected.address
         await super(EventHubConsumer, self)._open(timeout_time)
 
+    @_retry_decorator
+    async def _open_with_retry(self, timeout_time=None, **kwargs):
+        return await self._open(timeout_time=timeout_time, **kwargs)
+
     async def _receive(self, timeout_time=None, max_batch_size=None, **kwargs):
         last_exception = kwargs.get("last_exception")
         data_batch = kwargs.get("data_batch")
diff --git a/sdk/eventhub/azure-eventhubs/azure/eventhub/client.py b/sdk/eventhub/azure-eventhubs/azure/eventhub/client.py
index 1656ad27d074..0ee0a2bd8e62 100644
--- a/sdk/eventhub/azure-eventhubs/azure/eventhub/client.py
+++ b/sdk/eventhub/azure-eventhubs/azure/eventhub/client.py
@@ -108,10 +108,15 @@ def _close_connection(self):
         self._conn_manager.reset_connection_if_broken()
 
     def _management_request(self, mgmt_msg, op_type):
+        alt_creds = {
+            "username": self._auth_config.get("iot_username"),
+            "password": self._auth_config.get("iot_password")
+        }
+
         max_retries = self.config.max_retries
         retry_count = 0
         while retry_count <= self.config.max_retries:
-            mgmt_auth = self._create_auth()
+            mgmt_auth = self._create_auth(**alt_creds)
             mgmt_client = uamqp.AMQPClient(self.mgmt_target)
             try:
                 conn = self._conn_manager.get_connection(self.host, mgmt_auth)
diff --git a/sdk/eventhub/azure-eventhubs/azure/eventhub/consumer.py b/sdk/eventhub/azure-eventhubs/azure/eventhub/consumer.py
index 1f2b6db17728..84345fb97148 100644
--- a/sdk/eventhub/azure-eventhubs/azure/eventhub/consumer.py
+++ b/sdk/eventhub/azure-eventhubs/azure/eventhub/consumer.py
@@ -106,8 +106,10 @@ def __next__(self):
 
     def _create_handler(self):
         alt_creds = {
-            "username": self.client._auth_config.get("iot_username"),
-            "password": self.client._auth_config.get("iot_password")}
+            "username": self.client._auth_config.get("iot_username") if self.redirected else None,
+            "password": self.client._auth_config.get("iot_password") if self.redirected else None
+        }
+
         source = Source(self.source)
         if self.offset is not None:
             source.set_filter(self.offset._selector())
@@ -129,7 +131,7 @@ def _redirect(self, redirect):
         self.messages_iter = None
         super(EventHubConsumer, self)._redirect(redirect)
 
-    def _open(self, timeout_time=None):
+    def _open(self, timeout_time=None, **kwargs):
         """
         Open the EventHubConsumer using the supplied connection.
         If the handler has previously been redirected, the redirect
@@ -142,6 +144,10 @@ def _open(self, timeout_time=None):
             self.source = self.redirected.address
         super(EventHubConsumer, self)._open(timeout_time)
 
+    @_retry_decorator
+    def _open_with_retry(self, timeout_time=None, **kwargs):
+        return self._open(timeout_time=timeout_time, **kwargs)
+
     def _receive(self, timeout_time=None, max_batch_size=None, **kwargs):
         last_exception = kwargs.get("last_exception")
         data_batch = kwargs.get("data_batch")

From d37452c130da56f89df9e596fef7a3604db83942 Mon Sep 17 00:00:00 2001
From: Yunhao Ling <adam_ling@outlook.com>
Date: Tue, 20 Aug 2019 23:10:43 -0700
Subject: [PATCH 2/5] Support direct mgmt ops of iothub

---
 .../azure/eventhub/_consumer_producer_mixin.py      |  6 ++++--
 .../eventhub/aio/_consumer_producer_mixin_async.py  |  2 ++
 .../azure/eventhub/aio/client_async.py              | 13 +++++++++++++
 .../azure-eventhubs/azure/eventhub/client.py        | 13 +++++++++++++
 .../azure/eventhub/client_abstract.py               |  3 +++
 5 files changed, 35 insertions(+), 2 deletions(-)

diff --git a/sdk/eventhub/azure-eventhubs/azure/eventhub/_consumer_producer_mixin.py b/sdk/eventhub/azure-eventhubs/azure/eventhub/_consumer_producer_mixin.py
index 8e3cae1e8a0b..8ad49bcc4080 100644
--- a/sdk/eventhub/azure-eventhubs/azure/eventhub/_consumer_producer_mixin.py
+++ b/sdk/eventhub/azure-eventhubs/azure/eventhub/_consumer_producer_mixin.py
@@ -73,7 +73,7 @@ def _open(self, timeout_time=None):
             else:
                 alt_creds = {}
             self._create_handler()
-            self._handler.open(connection=self.client._conn_manager.get_connection(
+            self._handler.open(connection=self.client._conn_manager.get_connection(  # pylint: disable=protected-access
                 self.client.address.hostname,
                 self.client.get_auth(**alt_creds)
             ))
@@ -82,6 +82,8 @@ def _open(self, timeout_time=None):
             self._max_message_size_on_link = self._handler.message_handler._link.peer_max_message_size \
                                              or constants.MAX_MESSAGE_LENGTH_BYTES  # pylint: disable=protected-access
             self.running = True
+            if self.redirected and self.client._is_iothub:  # pylint: disable=protected-access
+                self.client._is_iothub_redirected = True  # pylint: disable=protected-access
 
     def _close_handler(self):
         self._handler.close()  # close the link (sharing connection) or connection (not sharing)
@@ -89,7 +91,7 @@ def _close_handler(self):
 
     def _close_connection(self):
         self._close_handler()
-        self.client._conn_manager.reset_connection_if_broken()
+        self.client._conn_manager.reset_connection_if_broken()  # pylint: disable=protected-access
 
     def _handle_exception(self, exception, retry_count, max_retries, timeout_time):
         if not self.running and isinstance(exception, compat.TimeoutException):
diff --git a/sdk/eventhub/azure-eventhubs/azure/eventhub/aio/_consumer_producer_mixin_async.py b/sdk/eventhub/azure-eventhubs/azure/eventhub/aio/_consumer_producer_mixin_async.py
index 8ccdcdddcd47..74028acf0bc7 100644
--- a/sdk/eventhub/azure-eventhubs/azure/eventhub/aio/_consumer_producer_mixin_async.py
+++ b/sdk/eventhub/azure-eventhubs/azure/eventhub/aio/_consumer_producer_mixin_async.py
@@ -83,6 +83,8 @@ async def _open(self, timeout_time=None):
             self._max_message_size_on_link = self._handler.message_handler._link.peer_max_message_size \
                                              or constants.MAX_MESSAGE_LENGTH_BYTES  # pylint: disable=protected-access
             self.running = True
+            if self.redirected and self.client._is_iothub:  # pylint: disable=protected-access
+                self.client._is_iothub_redirected = True  # pylint: disable=protected-access
 
     async def _close_handler(self):
         await self._handler.close_async()  # close the link (sharing connection) or connection (not sharing)
diff --git a/sdk/eventhub/azure-eventhubs/azure/eventhub/aio/client_async.py b/sdk/eventhub/azure-eventhubs/azure/eventhub/aio/client_async.py
index a6be575f43d8..ff6fee299d8f 100644
--- a/sdk/eventhub/azure-eventhubs/azure/eventhub/aio/client_async.py
+++ b/sdk/eventhub/azure-eventhubs/azure/eventhub/aio/client_async.py
@@ -126,6 +126,13 @@ async def _management_request(self, mgmt_msg, op_type):
             finally:
                 await mgmt_client.close_async()
 
+    async def _iothub_redirect(self):
+        async with self.create_consumer(consumer_group='$default',
+                                        partition_id='0',
+                                        event_position=EventPosition('-1'),
+                                        operation='/messages/events') as redirect_consumer:
+            await redirect_consumer._open_with_retry(timeout=self.config.receive_timeout)  # pylint: disable=protected-access
+
     async def get_properties(self):
         # type:() -> Dict[str, Any]
         """
@@ -139,6 +146,9 @@ async def get_properties(self):
         :rtype: dict
         :raises: ~azure.eventhub.ConnectError
         """
+        if self._is_iothub and not self._is_iothub_redirected:
+            await self._iothub_redirect()
+
         mgmt_msg = Message(application_properties={'name': self.eh_name})
         response = await self._management_request(mgmt_msg, op_type=b'com.microsoft:eventhub')
         output = {}
@@ -178,6 +188,9 @@ async def get_partition_properties(self, partition):
         :rtype: dict
         :raises: ~azure.eventhub.ConnectError
         """
+        if self._is_iothub and not self._is_iothub_redirected:
+            await self._iothub_redirect()
+
         mgmt_msg = Message(application_properties={'name': self.eh_name,
                                                    'partition': partition})
         response = await self._management_request(mgmt_msg, op_type=b'com.microsoft:partition')
diff --git a/sdk/eventhub/azure-eventhubs/azure/eventhub/client.py b/sdk/eventhub/azure-eventhubs/azure/eventhub/client.py
index 0ee0a2bd8e62..9e65453ba55d 100644
--- a/sdk/eventhub/azure-eventhubs/azure/eventhub/client.py
+++ b/sdk/eventhub/azure-eventhubs/azure/eventhub/client.py
@@ -134,6 +134,13 @@ def _management_request(self, mgmt_msg, op_type):
             finally:
                 mgmt_client.close()
 
+    def _iothub_redirect(self):
+        with self.create_consumer(consumer_group='$default',
+                                  partition_id='0',
+                                  event_position=EventPosition('-1'),
+                                  operation='/messages/events') as redirect_consumer:
+            redirect_consumer._open_with_retry(timeout=self.config.receive_timeout)  # pylint: disable=protected-access
+
     def get_properties(self):
         # type:() -> Dict[str, Any]
         """
@@ -147,6 +154,9 @@ def get_properties(self):
         :rtype: dict
         :raises: ~azure.eventhub.ConnectError
         """
+        if self._is_iothub and not self._is_iothub_redirected:
+            self._iothub_redirect()
+
         mgmt_msg = Message(application_properties={'name': self.eh_name})
         response = self._management_request(mgmt_msg, op_type=b'com.microsoft:eventhub')
         output = {}
@@ -186,6 +196,9 @@ def get_partition_properties(self, partition):
         :rtype: dict
         :raises: ~azure.eventhub.ConnectError
         """
+        if self._is_iothub and not self._is_iothub_redirected:
+            self._iothub_redirect()
+
         mgmt_msg = Message(application_properties={'name': self.eh_name,
                                                    'partition': partition})
         response = self._management_request(mgmt_msg, op_type=b'com.microsoft:partition')
diff --git a/sdk/eventhub/azure-eventhubs/azure/eventhub/client_abstract.py b/sdk/eventhub/azure-eventhubs/azure/eventhub/client_abstract.py
index dd62e3e76de1..a6433bfbdc9e 100644
--- a/sdk/eventhub/azure-eventhubs/azure/eventhub/client_abstract.py
+++ b/sdk/eventhub/azure-eventhubs/azure/eventhub/client_abstract.py
@@ -157,6 +157,8 @@ def __init__(self, host, event_hub_path, credential, **kwargs):
         self.get_auth = functools.partial(self._create_auth)
         self.config = _Configuration(**kwargs)
         self.debug = self.config.network_tracing
+        self._is_iothub = kwargs.get("is_iothub", False)
+        self._is_iothub_redirected = False
 
         log.info("%r: Created the Event Hub client", self.container_id)
 
@@ -279,6 +281,7 @@ def from_connection_string(cls, conn_str, **kwargs):
             kwargs.pop("event_hub_path", None)
             return cls(host, entity, EventHubSharedKeyCredential(policy, key), **kwargs)
         else:
+            kwargs['is_iothub'] = True
             return cls._from_iothub_connection_string(conn_str, **kwargs)
 
     @abstractmethod

From b5e9b092485b9b1da0b89423bbeaa68cc51a2acf Mon Sep 17 00:00:00 2001
From: Yunhao Ling <adam_ling@outlook.com>
Date: Wed, 21 Aug 2019 17:37:36 -0700
Subject: [PATCH 3/5] Improve mgmt ops and update livetest

---
 .../eventhub/_consumer_producer_mixin.py      |  2 -
 .../aio/_consumer_producer_mixin_async.py     |  2 -
 .../azure/eventhub/aio/client_async.py        | 22 +++++-
 .../azure/eventhub/aio/consumer_async.py      |  8 ++-
 .../azure/eventhub/aio/producer_async.py      |  2 +-
 .../azure-eventhubs/azure/eventhub/client.py  | 18 ++++-
 .../azure/eventhub/client_abstract.py         | 11 +--
 .../azure/eventhub/consumer.py                |  3 +
 .../asynctests/test_iothub_receive_async.py   | 67 +++++++++++++++----
 .../tests/test_iothub_receive.py              | 53 +++++++++++++--
 10 files changed, 146 insertions(+), 42 deletions(-)

diff --git a/sdk/eventhub/azure-eventhubs/azure/eventhub/_consumer_producer_mixin.py b/sdk/eventhub/azure-eventhubs/azure/eventhub/_consumer_producer_mixin.py
index 8ad49bcc4080..208197c956f0 100644
--- a/sdk/eventhub/azure-eventhubs/azure/eventhub/_consumer_producer_mixin.py
+++ b/sdk/eventhub/azure-eventhubs/azure/eventhub/_consumer_producer_mixin.py
@@ -82,8 +82,6 @@ def _open(self, timeout_time=None):
             self._max_message_size_on_link = self._handler.message_handler._link.peer_max_message_size \
                                              or constants.MAX_MESSAGE_LENGTH_BYTES  # pylint: disable=protected-access
             self.running = True
-            if self.redirected and self.client._is_iothub:  # pylint: disable=protected-access
-                self.client._is_iothub_redirected = True  # pylint: disable=protected-access
 
     def _close_handler(self):
         self._handler.close()  # close the link (sharing connection) or connection (not sharing)
diff --git a/sdk/eventhub/azure-eventhubs/azure/eventhub/aio/_consumer_producer_mixin_async.py b/sdk/eventhub/azure-eventhubs/azure/eventhub/aio/_consumer_producer_mixin_async.py
index 74028acf0bc7..8ccdcdddcd47 100644
--- a/sdk/eventhub/azure-eventhubs/azure/eventhub/aio/_consumer_producer_mixin_async.py
+++ b/sdk/eventhub/azure-eventhubs/azure/eventhub/aio/_consumer_producer_mixin_async.py
@@ -83,8 +83,6 @@ async def _open(self, timeout_time=None):
             self._max_message_size_on_link = self._handler.message_handler._link.peer_max_message_size \
                                              or constants.MAX_MESSAGE_LENGTH_BYTES  # pylint: disable=protected-access
             self.running = True
-            if self.redirected and self.client._is_iothub:  # pylint: disable=protected-access
-                self.client._is_iothub_redirected = True  # pylint: disable=protected-access
 
     async def _close_handler(self):
         await self._handler.close_async()  # close the link (sharing connection) or connection (not sharing)
diff --git a/sdk/eventhub/azure-eventhubs/azure/eventhub/aio/client_async.py b/sdk/eventhub/azure-eventhubs/azure/eventhub/aio/client_async.py
index ff6fee299d8f..2be245e88390 100644
--- a/sdk/eventhub/azure-eventhubs/azure/eventhub/aio/client_async.py
+++ b/sdk/eventhub/azure-eventhubs/azure/eventhub/aio/client_async.py
@@ -6,6 +6,11 @@
 import datetime
 import functools
 import asyncio
+try:
+    from urlparse import urlparse
+    from urllib import unquote_plus, urlencode, quote_plus
+except ImportError:
+    from urllib.parse import urlparse, unquote_plus, urlencode, quote_plus
 from typing import Any, List, Dict
 
 from uamqp import authentication, constants
@@ -43,6 +48,7 @@ class EventHubClient(EventHubClientAbstract):
 
     def __init__(self, host, event_hub_path, credential, **kwargs):
         super(EventHubClient, self).__init__(host, event_hub_path, credential, **kwargs)
+        self._lock = asyncio.Lock()
         self._conn_manager = get_connection_manager(**kwargs)
 
     async def __aenter__(self):
@@ -133,6 +139,18 @@ async def _iothub_redirect(self):
                                         operation='/messages/events') as redirect_consumer:
             await redirect_consumer._open_with_retry(timeout=self.config.receive_timeout)  # pylint: disable=protected-access
 
+    async def _process_redirect_uri(self, redirect):
+        async with self._lock:
+            redirect_uri = redirect.address.decode('utf-8')
+            auth_uri, _, _ = redirect_uri.partition("/ConsumerGroups")
+            self.address = urlparse(auth_uri)
+            self.host = self.address.hostname
+            self.auth_uri = "sb://{}{}".format(self.address.hostname, self.address.path)
+            self.eh_name = self.address.path.lstrip('/')
+            self.mgmt_target = redirect_uri
+            if self._is_iothub:
+                self._iothub_redirected = redirect
+
     async def get_properties(self):
         # type:() -> Dict[str, Any]
         """
@@ -146,7 +164,7 @@ async def get_properties(self):
         :rtype: dict
         :raises: ~azure.eventhub.ConnectError
         """
-        if self._is_iothub and not self._is_iothub_redirected:
+        if self._is_iothub and not self._iothub_redirected:
             await self._iothub_redirect()
 
         mgmt_msg = Message(application_properties={'name': self.eh_name})
@@ -188,7 +206,7 @@ async def get_partition_properties(self, partition):
         :rtype: dict
         :raises: ~azure.eventhub.ConnectError
         """
-        if self._is_iothub and not self._is_iothub_redirected:
+        if self._is_iothub and not self._iothub_redirected:
             await self._iothub_redirect()
 
         mgmt_msg = Message(application_properties={'name': self.eh_name,
diff --git a/sdk/eventhub/azure-eventhubs/azure/eventhub/aio/consumer_async.py b/sdk/eventhub/azure-eventhubs/azure/eventhub/aio/consumer_async.py
index 6ca674320071..cb818a31790a 100644
--- a/sdk/eventhub/azure-eventhubs/azure/eventhub/aio/consumer_async.py
+++ b/sdk/eventhub/azure-eventhubs/azure/eventhub/aio/consumer_async.py
@@ -143,8 +143,11 @@ async def _open(self, timeout_time=None, **kwargs):
 
         """
         # pylint: disable=protected-access
+        if self.client._iothub_redirected:
+            self.redirected = self.client._iothub_redirected
+
         if not self.running and self.redirected:
-            self.client._process_redirect_uri(self.redirected)
+            await self.client._process_redirect_uri(self.redirected)
             self.source = self.redirected.address
         await super(EventHubConsumer, self)._open(timeout_time)
 
@@ -259,4 +262,5 @@ async def close(self, exception=None):
             self.error = EventHubError(str(exception))
         else:
             self.error = EventHubError("This receive handler is now closed.")
-        await self._handler.close_async()
+        if self._handler:
+            await self._handler.close_async()
diff --git a/sdk/eventhub/azure-eventhubs/azure/eventhub/aio/producer_async.py b/sdk/eventhub/azure-eventhubs/azure/eventhub/aio/producer_async.py
index c45cf9a6b283..bdd79bb72406 100644
--- a/sdk/eventhub/azure-eventhubs/azure/eventhub/aio/producer_async.py
+++ b/sdk/eventhub/azure-eventhubs/azure/eventhub/aio/producer_async.py
@@ -106,7 +106,7 @@ async def _open(self, timeout_time=None, **kwargs):
 
         """
         if not self.running and self.redirected:
-            self.client._process_redirect_uri(self.redirected)
+            await self.client._process_redirect_uri(self.redirected)
             self.target = self.redirected.address
         await super(EventHubProducer, self)._open(timeout_time)
 
diff --git a/sdk/eventhub/azure-eventhubs/azure/eventhub/client.py b/sdk/eventhub/azure-eventhubs/azure/eventhub/client.py
index 9e65453ba55d..c5b76ca44317 100644
--- a/sdk/eventhub/azure-eventhubs/azure/eventhub/client.py
+++ b/sdk/eventhub/azure-eventhubs/azure/eventhub/client.py
@@ -7,6 +7,7 @@
 import logging
 import datetime
 import functools
+import threading
 try:
     from urlparse import urlparse
     from urllib import unquote_plus, urlencode, quote_plus
@@ -48,6 +49,7 @@ class EventHubClient(EventHubClientAbstract):
 
     def __init__(self, host, event_hub_path, credential, **kwargs):
         super(EventHubClient, self).__init__(host, event_hub_path, credential, **kwargs)
+        self._lock = threading.Lock()
         self._conn_manager = get_connection_manager(**kwargs)
 
     def __enter__(self):
@@ -141,6 +143,18 @@ def _iothub_redirect(self):
                                   operation='/messages/events') as redirect_consumer:
             redirect_consumer._open_with_retry(timeout=self.config.receive_timeout)  # pylint: disable=protected-access
 
+    def _process_redirect_uri(self, redirect):
+        with self._lock:
+            redirect_uri = redirect.address.decode('utf-8')
+            auth_uri, _, _ = redirect_uri.partition("/ConsumerGroups")
+            self.address = urlparse(auth_uri)
+            self.host = self.address.hostname
+            self.auth_uri = "sb://{}{}".format(self.address.hostname, self.address.path)
+            self.eh_name = self.address.path.lstrip('/')
+            self.mgmt_target = redirect_uri
+            if self._is_iothub:
+                self._iothub_redirected = redirect
+
     def get_properties(self):
         # type:() -> Dict[str, Any]
         """
@@ -154,7 +168,7 @@ def get_properties(self):
         :rtype: dict
         :raises: ~azure.eventhub.ConnectError
         """
-        if self._is_iothub and not self._is_iothub_redirected:
+        if self._is_iothub and not self._iothub_redirected:
             self._iothub_redirect()
 
         mgmt_msg = Message(application_properties={'name': self.eh_name})
@@ -196,7 +210,7 @@ def get_partition_properties(self, partition):
         :rtype: dict
         :raises: ~azure.eventhub.ConnectError
         """
-        if self._is_iothub and not self._is_iothub_redirected:
+        if self._is_iothub and not self._iothub_redirected:
             self._iothub_redirect()
 
         mgmt_msg = Message(application_properties={'name': self.eh_name,
diff --git a/sdk/eventhub/azure-eventhubs/azure/eventhub/client_abstract.py b/sdk/eventhub/azure-eventhubs/azure/eventhub/client_abstract.py
index a6433bfbdc9e..91a5aaf211f8 100644
--- a/sdk/eventhub/azure-eventhubs/azure/eventhub/client_abstract.py
+++ b/sdk/eventhub/azure-eventhubs/azure/eventhub/client_abstract.py
@@ -158,7 +158,7 @@ def __init__(self, host, event_hub_path, credential, **kwargs):
         self.config = _Configuration(**kwargs)
         self.debug = self.config.network_tracing
         self._is_iothub = kwargs.get("is_iothub", False)
-        self._is_iothub_redirected = False
+        self._iothub_redirected = None
 
         log.info("%r: Created the Event Hub client", self.container_id)
 
@@ -211,15 +211,6 @@ def _create_properties(self, user_agent=None):  # pylint: disable=no-self-use
         properties["user-agent"] = final_user_agent
         return properties
 
-    def _process_redirect_uri(self, redirect):
-        redirect_uri = redirect.address.decode('utf-8')
-        auth_uri, _, _ = redirect_uri.partition("/ConsumerGroups")
-        self.address = urlparse(auth_uri)
-        self.host = self.address.hostname
-        self.auth_uri = "sb://{}{}".format(self.address.hostname, self.address.path)
-        self.eh_name = self.address.path.lstrip('/')
-        self.mgmt_target = redirect_uri
-
     @classmethod
     def from_connection_string(cls, conn_str, **kwargs):
         """Create an EventHubClient from an EventHub/IotHub connection string.
diff --git a/sdk/eventhub/azure-eventhubs/azure/eventhub/consumer.py b/sdk/eventhub/azure-eventhubs/azure/eventhub/consumer.py
index 84345fb97148..eab8990eb40d 100644
--- a/sdk/eventhub/azure-eventhubs/azure/eventhub/consumer.py
+++ b/sdk/eventhub/azure-eventhubs/azure/eventhub/consumer.py
@@ -139,6 +139,9 @@ def _open(self, timeout_time=None, **kwargs):
 
         """
         # pylint: disable=protected-access
+        if self.client._iothub_redirected:
+            self.redirected = self.client._iothub_redirected
+
         if not self.running and self.redirected:
             self.client._process_redirect_uri(self.redirected)
             self.source = self.redirected.address
diff --git a/sdk/eventhub/azure-eventhubs/tests/asynctests/test_iothub_receive_async.py b/sdk/eventhub/azure-eventhubs/tests/asynctests/test_iothub_receive_async.py
index f581f64584ab..9226a5d62a93 100644
--- a/sdk/eventhub/azure-eventhubs/tests/asynctests/test_iothub_receive_async.py
+++ b/sdk/eventhub/azure-eventhubs/tests/asynctests/test_iothub_receive_async.py
@@ -4,13 +4,11 @@
 # license information.
 #--------------------------------------------------------------------------
 
-import os
 import asyncio
 import pytest
-import time
 
 from azure.eventhub.aio import EventHubClient
-from azure.eventhub import EventData, EventPosition, EventHubError
+from azure.eventhub import EventPosition
 
 
 async def pump(receiver, sleep=None):
@@ -18,25 +16,16 @@ async def pump(receiver, sleep=None):
     if sleep:
         await asyncio.sleep(sleep)
     async with receiver:
-        batch = await receiver.receive(timeout=1)
+        batch = await receiver.receive(timeout=3)
         messages += len(batch)
     return messages
 
 
-async def get_partitions(iot_connection_str):
-    client = EventHubClient.from_connection_string(iot_connection_str, network_tracing=False)
-    receiver = client.create_consumer(consumer_group="$default", partition_id="0", event_position=EventPosition("-1"), prefetch=1000, operation='/messages/events')
-    async with receiver:
-        partitions = await client.get_properties()
-        return partitions["partition_ids"]
-
-
 @pytest.mark.liveTest
 @pytest.mark.asyncio
 async def test_iothub_receive_multiple_async(iot_connection_str):
-    pytest.skip("This will get AuthenticationError. We're investigating...")
-    partitions = await get_partitions(iot_connection_str)
     client = EventHubClient.from_connection_string(iot_connection_str, network_tracing=False)
+    partitions = await client.get_partition_ids()
     receivers = []
     for p in partitions:
         receivers.append(client.create_consumer(consumer_group="$default", partition_id=p, event_position=EventPosition("-1"), prefetch=10, operation='/messages/events'))
@@ -44,3 +33,53 @@ async def test_iothub_receive_multiple_async(iot_connection_str):
 
     assert isinstance(outputs[0], int) and outputs[0] <= 10
     assert isinstance(outputs[1], int) and outputs[1] <= 10
+
+
+@pytest.mark.liveTest
+@pytest.mark.asyncio
+async def test_iothub_get_properties_async(iot_connection_str, device_id):
+    client = EventHubClient.from_connection_string(iot_connection_str, network_tracing=False)
+    properties = await client.get_properties()
+    assert properties["partition_ids"] == ["0", "1", "2", "3"]
+
+
+@pytest.mark.liveTest
+@pytest.mark.asyncio
+async def test_iothub_get_partition_ids_async(iot_connection_str, device_id):
+    client = EventHubClient.from_connection_string(iot_connection_str, network_tracing=False)
+    partitions = await client.get_partition_ids()
+    assert partitions == ["0", "1", "2", "3"]
+
+
+@pytest.mark.liveTest
+@pytest.mark.asyncio
+async def test_iothub_get_partition_properties_async(iot_connection_str, device_id):
+    client = EventHubClient.from_connection_string(iot_connection_str, network_tracing=False)
+    partition_properties = await client.get_partition_properties("0")
+    assert partition_properties["id"] == "0"
+
+
+@pytest.mark.liveTest
+@pytest.mark.asyncio
+async def test_iothub_receive_after_mgmt_ops_async(iot_connection_str, device_id):
+    client = EventHubClient.from_connection_string(iot_connection_str, network_tracing=False)
+    partitions = await client.get_partition_ids()
+    assert partitions == ["0", "1", "2", "3"]
+    receiver = client.create_consumer(consumer_group="$default", partition_id=partitions[0], event_position=EventPosition("-1"), operation='/messages/events')
+    async with receiver:
+        received = await receiver.receive(timeout=5)
+        assert len(received) == 0
+
+
+@pytest.mark.liveTest
+@pytest.mark.asyncio
+async def test_iothub_mgmt_ops_after_receive_async(iot_connection_str, device_id):
+    client = EventHubClient.from_connection_string(iot_connection_str, network_tracing=False)
+    receiver = client.create_consumer(consumer_group="$default", partition_id="0", event_position=EventPosition("-1"), operation='/messages/events')
+    async with receiver:
+        received = await receiver.receive(timeout=5)
+        assert len(received) == 0
+
+    partitions = await client.get_partition_ids()
+    assert partitions == ["0", "1", "2", "3"]
+
diff --git a/sdk/eventhub/azure-eventhubs/tests/test_iothub_receive.py b/sdk/eventhub/azure-eventhubs/tests/test_iothub_receive.py
index ac5787b6b12e..680b7fbc3a58 100644
--- a/sdk/eventhub/azure-eventhubs/tests/test_iothub_receive.py
+++ b/sdk/eventhub/azure-eventhubs/tests/test_iothub_receive.py
@@ -4,23 +4,62 @@
 # license information.
 #--------------------------------------------------------------------------
 
-import os
 import pytest
-import time
 
-from azure.eventhub import EventData, EventPosition, EventHubClient
+from azure.eventhub import EventPosition, EventHubClient
 
 
 @pytest.mark.liveTest
 def test_iothub_receive_sync(iot_connection_str, device_id):
-    pytest.skip("current code will cause ErrorCodes.LinkRedirect")
+    print("!!", iot_connection_str)
     client = EventHubClient.from_connection_string(iot_connection_str, network_tracing=False)
     receiver = client.create_consumer(consumer_group="$default", partition_id="0", event_position=EventPosition("-1"), operation='/messages/events')
-    receiver._open()
     try:
-        partitions = client.get_properties()
-        assert partitions["partition_ids"] == ["0", "1", "2", "3"]
         received = receiver.receive(timeout=5)
         assert len(received) == 0
     finally:
         receiver.close()
+
+
+@pytest.mark.liveTest
+def test_iothub_get_properties_sync(iot_connection_str, device_id):
+    client = EventHubClient.from_connection_string(iot_connection_str, network_tracing=False)
+    properties = client.get_properties()
+    assert properties["partition_ids"] == ["0", "1", "2", "3"]
+
+
+@pytest.mark.liveTest
+def test_iothub_get_partition_ids_sync(iot_connection_str, device_id):
+    client = EventHubClient.from_connection_string(iot_connection_str, network_tracing=False)
+    partitions = client.get_partition_ids()
+    assert partitions == ["0", "1", "2", "3"]
+
+
+@pytest.mark.liveTest
+def test_iothub_get_partition_properties_sync(iot_connection_str, device_id):
+    client = EventHubClient.from_connection_string(iot_connection_str, network_tracing=False)
+    partition_properties = client.get_partition_properties("0")
+    assert partition_properties["id"] == "0"
+
+
+@pytest.mark.liveTest
+def test_iothub_receive_after_mgmt_ops_sync(iot_connection_str, device_id):
+    client = EventHubClient.from_connection_string(iot_connection_str, network_tracing=False)
+    partitions = client.get_partition_ids()
+    assert partitions == ["0", "1", "2", "3"]
+    receiver = client.create_consumer(consumer_group="$default", partition_id=partitions[0], event_position=EventPosition("-1"), operation='/messages/events')
+    with receiver:
+        received = receiver.receive(timeout=5)
+        assert len(received) == 0
+
+
+@pytest.mark.liveTest
+def test_iothub_mgmt_ops_after_receive_sync(iot_connection_str, device_id):
+    client = EventHubClient.from_connection_string(iot_connection_str, network_tracing=False)
+    receiver = client.create_consumer(consumer_group="$default", partition_id="0", event_position=EventPosition("-1"), operation='/messages/events')
+    with receiver:
+        received = receiver.receive(timeout=5)
+        assert len(received) == 0
+
+    partitions = client.get_partition_ids()
+    assert partitions == ["0", "1", "2", "3"]

From ba31bb01e920cc2e6995006ed9a6be9b431bd843 Mon Sep 17 00:00:00 2001
From: Yunhao Ling <adam_ling@outlook.com>
Date: Wed, 21 Aug 2019 17:39:00 -0700
Subject: [PATCH 4/5] Small fix

---
 sdk/eventhub/azure-eventhubs/tests/test_iothub_receive.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/sdk/eventhub/azure-eventhubs/tests/test_iothub_receive.py b/sdk/eventhub/azure-eventhubs/tests/test_iothub_receive.py
index 680b7fbc3a58..ac7e211bd736 100644
--- a/sdk/eventhub/azure-eventhubs/tests/test_iothub_receive.py
+++ b/sdk/eventhub/azure-eventhubs/tests/test_iothub_receive.py
@@ -11,7 +11,6 @@
 
 @pytest.mark.liveTest
 def test_iothub_receive_sync(iot_connection_str, device_id):
-    print("!!", iot_connection_str)
     client = EventHubClient.from_connection_string(iot_connection_str, network_tracing=False)
     receiver = client.create_consumer(consumer_group="$default", partition_id="0", event_position=EventPosition("-1"), operation='/messages/events')
     try:

From afd8df6c8def2d77aee664d8a502975480018623 Mon Sep 17 00:00:00 2001
From: Yunhao Ling <adam_ling@outlook.com>
Date: Wed, 28 Aug 2019 00:17:49 -0700
Subject: [PATCH 5/5] Improvement of iothub mgmt

---
 .../azure/eventhub/aio/client_async.py        | 34 ++++++++-----------
 .../azure/eventhub/aio/consumer_async.py      |  5 ++-
 .../azure-eventhubs/azure/eventhub/client.py  | 33 +++++++-----------
 .../azure/eventhub/client_abstract.py         | 23 ++++++++++---
 .../azure/eventhub/consumer.py                |  3 +-
 5 files changed, 48 insertions(+), 50 deletions(-)

diff --git a/sdk/eventhub/azure-eventhubs/azure/eventhub/aio/client_async.py b/sdk/eventhub/azure-eventhubs/azure/eventhub/aio/client_async.py
index 2be245e88390..e5227ca9b585 100644
--- a/sdk/eventhub/azure-eventhubs/azure/eventhub/aio/client_async.py
+++ b/sdk/eventhub/azure-eventhubs/azure/eventhub/aio/client_async.py
@@ -107,6 +107,9 @@ async def _close_connection(self):
         await self._conn_manager.reset_connection_if_broken()
 
     async def _management_request(self, mgmt_msg, op_type):
+        if self._is_iothub and not self._iothub_redirect_info:
+            await self._iothub_redirect()
+
         alt_creds = {
             "username": self._auth_config.get("iot_username"),
             "password": self._auth_config.get("iot_password")
@@ -133,23 +136,16 @@ async def _management_request(self, mgmt_msg, op_type):
                 await mgmt_client.close_async()
 
     async def _iothub_redirect(self):
-        async with self.create_consumer(consumer_group='$default',
-                                        partition_id='0',
-                                        event_position=EventPosition('-1'),
-                                        operation='/messages/events') as redirect_consumer:
-            await redirect_consumer._open_with_retry(timeout=self.config.receive_timeout)  # pylint: disable=protected-access
-
-    async def _process_redirect_uri(self, redirect):
         async with self._lock:
-            redirect_uri = redirect.address.decode('utf-8')
-            auth_uri, _, _ = redirect_uri.partition("/ConsumerGroups")
-            self.address = urlparse(auth_uri)
-            self.host = self.address.hostname
-            self.auth_uri = "sb://{}{}".format(self.address.hostname, self.address.path)
-            self.eh_name = self.address.path.lstrip('/')
-            self.mgmt_target = redirect_uri
-            if self._is_iothub:
-                self._iothub_redirected = redirect
+            if self._is_iothub and not self._iothub_redirect_info:
+                if not self._redirect_consumer:
+                    self._redirect_consumer = self.create_consumer(consumer_group='$default',
+                                                                   partition_id='0',
+                                                                   event_position=EventPosition('-1'),
+                                                                   operation='/messages/events')
+                async with self._redirect_consumer:
+                    await self._redirect_consumer._open_with_retry(timeout=self.config.receive_timeout)  # pylint: disable=protected-access
+                self._redirect_consumer = None
 
     async def get_properties(self):
         # type:() -> Dict[str, Any]
@@ -164,9 +160,8 @@ async def get_properties(self):
         :rtype: dict
         :raises: ~azure.eventhub.ConnectError
         """
-        if self._is_iothub and not self._iothub_redirected:
+        if self._is_iothub and not self._iothub_redirect_info:
             await self._iothub_redirect()
-
         mgmt_msg = Message(application_properties={'name': self.eh_name})
         response = await self._management_request(mgmt_msg, op_type=b'com.microsoft:eventhub')
         output = {}
@@ -206,9 +201,8 @@ async def get_partition_properties(self, partition):
         :rtype: dict
         :raises: ~azure.eventhub.ConnectError
         """
-        if self._is_iothub and not self._iothub_redirected:
+        if self._is_iothub and not self._iothub_redirect_info:
             await self._iothub_redirect()
-
         mgmt_msg = Message(application_properties={'name': self.eh_name,
                                                    'partition': partition})
         response = await self._management_request(mgmt_msg, op_type=b'com.microsoft:partition')
diff --git a/sdk/eventhub/azure-eventhubs/azure/eventhub/aio/consumer_async.py b/sdk/eventhub/azure-eventhubs/azure/eventhub/aio/consumer_async.py
index cb818a31790a..0b604a37688f 100644
--- a/sdk/eventhub/azure-eventhubs/azure/eventhub/aio/consumer_async.py
+++ b/sdk/eventhub/azure-eventhubs/azure/eventhub/aio/consumer_async.py
@@ -143,11 +143,10 @@ async def _open(self, timeout_time=None, **kwargs):
 
         """
         # pylint: disable=protected-access
-        if self.client._iothub_redirected:
-            self.redirected = self.client._iothub_redirected
+        self.redirected = self.redirected or self.client._iothub_redirect_info
 
         if not self.running and self.redirected:
-            await self.client._process_redirect_uri(self.redirected)
+            self.client._process_redirect_uri(self.redirected)
             self.source = self.redirected.address
         await super(EventHubConsumer, self)._open(timeout_time)
 
diff --git a/sdk/eventhub/azure-eventhubs/azure/eventhub/client.py b/sdk/eventhub/azure-eventhubs/azure/eventhub/client.py
index c5b76ca44317..537b905af2f6 100644
--- a/sdk/eventhub/azure-eventhubs/azure/eventhub/client.py
+++ b/sdk/eventhub/azure-eventhubs/azure/eventhub/client.py
@@ -49,7 +49,7 @@ class EventHubClient(EventHubClientAbstract):
 
     def __init__(self, host, event_hub_path, credential, **kwargs):
         super(EventHubClient, self).__init__(host, event_hub_path, credential, **kwargs)
-        self._lock = threading.Lock()
+        self._lock = threading.RLock()
         self._conn_manager = get_connection_manager(**kwargs)
 
     def __enter__(self):
@@ -137,23 +137,16 @@ def _management_request(self, mgmt_msg, op_type):
                 mgmt_client.close()
 
     def _iothub_redirect(self):
-        with self.create_consumer(consumer_group='$default',
-                                  partition_id='0',
-                                  event_position=EventPosition('-1'),
-                                  operation='/messages/events') as redirect_consumer:
-            redirect_consumer._open_with_retry(timeout=self.config.receive_timeout)  # pylint: disable=protected-access
-
-    def _process_redirect_uri(self, redirect):
         with self._lock:
-            redirect_uri = redirect.address.decode('utf-8')
-            auth_uri, _, _ = redirect_uri.partition("/ConsumerGroups")
-            self.address = urlparse(auth_uri)
-            self.host = self.address.hostname
-            self.auth_uri = "sb://{}{}".format(self.address.hostname, self.address.path)
-            self.eh_name = self.address.path.lstrip('/')
-            self.mgmt_target = redirect_uri
-            if self._is_iothub:
-                self._iothub_redirected = redirect
+            if self._is_iothub and not self._iothub_redirect_info:
+                if not self._redirect_consumer:
+                    self._redirect_consumer = self.create_consumer(consumer_group='$default',
+                                                                   partition_id='0',
+                                                                   event_position=EventPosition('-1'),
+                                                                   operation='/messages/events')
+                with self._redirect_consumer:
+                    self._redirect_consumer._open_with_retry(timeout=self.config.receive_timeout)  # pylint: disable=protected-access
+                self._redirect_consumer = None
 
     def get_properties(self):
         # type:() -> Dict[str, Any]
@@ -168,9 +161,8 @@ def get_properties(self):
         :rtype: dict
         :raises: ~azure.eventhub.ConnectError
         """
-        if self._is_iothub and not self._iothub_redirected:
+        if self._is_iothub and not self._iothub_redirect_info:
             self._iothub_redirect()
-
         mgmt_msg = Message(application_properties={'name': self.eh_name})
         response = self._management_request(mgmt_msg, op_type=b'com.microsoft:eventhub')
         output = {}
@@ -210,9 +202,8 @@ def get_partition_properties(self, partition):
         :rtype: dict
         :raises: ~azure.eventhub.ConnectError
         """
-        if self._is_iothub and not self._iothub_redirected:
+        if self._is_iothub and not self._iothub_redirect_info:
             self._iothub_redirect()
-
         mgmt_msg = Message(application_properties={'name': self.eh_name,
                                                    'partition': partition})
         response = self._management_request(mgmt_msg, op_type=b'com.microsoft:partition')
diff --git a/sdk/eventhub/azure-eventhubs/azure/eventhub/client_abstract.py b/sdk/eventhub/azure-eventhubs/azure/eventhub/client_abstract.py
index 91a5aaf211f8..072b84179fc7 100644
--- a/sdk/eventhub/azure-eventhubs/azure/eventhub/client_abstract.py
+++ b/sdk/eventhub/azure-eventhubs/azure/eventhub/client_abstract.py
@@ -24,7 +24,7 @@
     from azure.core.credentials import TokenCredential
     from typing import Union, Any
 
-from azure.eventhub import __version__
+from azure.eventhub import __version__, EventPosition
 from azure.eventhub.configuration import _Configuration
 from .common import EventHubSharedKeyCredential, EventHubSASTokenCredential, _Address
 
@@ -157,8 +157,8 @@ def __init__(self, host, event_hub_path, credential, **kwargs):
         self.get_auth = functools.partial(self._create_auth)
         self.config = _Configuration(**kwargs)
         self.debug = self.config.network_tracing
-        self._is_iothub = kwargs.get("is_iothub", False)
-        self._iothub_redirected = None
+        self._is_iothub = False
+        self._iothub_redirect_info = None
 
         log.info("%r: Created the Event Hub client", self.container_id)
 
@@ -179,6 +179,11 @@ def _from_iothub_connection_string(cls, conn_str, **kwargs):
             'iot_password': key,
             'username': username,
             'password': password}
+        client._is_iothub = True
+        client._redirect_consumer = client.create_consumer(consumer_group='$default',
+                                                           partition_id='0',
+                                                           event_position=EventPosition('-1'),
+                                                           operation='/messages/events')
         return client
 
     @abstractmethod
@@ -211,6 +216,17 @@ def _create_properties(self, user_agent=None):  # pylint: disable=no-self-use
         properties["user-agent"] = final_user_agent
         return properties
 
+    def _process_redirect_uri(self, redirect):
+        redirect_uri = redirect.address.decode('utf-8')
+        auth_uri, _, _ = redirect_uri.partition("/ConsumerGroups")
+        self.address = urlparse(auth_uri)
+        self.host = self.address.hostname
+        self.auth_uri = "sb://{}{}".format(self.address.hostname, self.address.path)
+        self.eh_name = self.address.path.lstrip('/')
+        self.mgmt_target = redirect_uri
+        if self._is_iothub:
+            self._iothub_redirect_info = redirect
+
     @classmethod
     def from_connection_string(cls, conn_str, **kwargs):
         """Create an EventHubClient from an EventHub/IotHub connection string.
@@ -272,7 +288,6 @@ def from_connection_string(cls, conn_str, **kwargs):
             kwargs.pop("event_hub_path", None)
             return cls(host, entity, EventHubSharedKeyCredential(policy, key), **kwargs)
         else:
-            kwargs['is_iothub'] = True
             return cls._from_iothub_connection_string(conn_str, **kwargs)
 
     @abstractmethod
diff --git a/sdk/eventhub/azure-eventhubs/azure/eventhub/consumer.py b/sdk/eventhub/azure-eventhubs/azure/eventhub/consumer.py
index eab8990eb40d..930570d51e02 100644
--- a/sdk/eventhub/azure-eventhubs/azure/eventhub/consumer.py
+++ b/sdk/eventhub/azure-eventhubs/azure/eventhub/consumer.py
@@ -139,8 +139,7 @@ def _open(self, timeout_time=None, **kwargs):
 
         """
         # pylint: disable=protected-access
-        if self.client._iothub_redirected:
-            self.redirected = self.client._iothub_redirected
+        self.redirected = self.redirected or self.client._iothub_redirect_info
 
         if not self.running and self.redirected:
             self.client._process_redirect_uri(self.redirected)