e2e: replace timing waits with bounded synchronization

dev_utils/dev_utils/service_helper.py

@@ -29,6 +29,9 @@ async def _run_in_executor(self, func, *args):
     def set_identity(self, device_id, module_id):
         return self._inner_object.set_identity(device_id, module_id)
 
+    async def wait_until_ready(self, timeout):
+        return await self._run_in_executor(self._inner_object.wait_until_ready, timeout)
+
     async def set_desired_properties(self, desired_props):
         return await self._run_in_executor(self._inner_object.set_desired_properties, desired_props)
 
@@ -54,11 +57,12 @@ async def send_c2d(
     ):
         return await self._run_in_executor(self._inner_object.send_c2d, payload, properties)
 
-    async def wait_for_eventhub_arrival(self, message_id, timeout=60):
+    async def wait_for_eventhub_arrival(self, message_id, timeout=60, event_filter=None):
         return await self._run_in_executor(
             self._inner_object.wait_for_eventhub_arrival,
             message_id,
             timeout,
+            event_filter,
         )
 
     async def get_next_reported_patch_arrival(self, block=True, timeout=240):

dev_utils/dev_utils/service_helper_sync.py

@@ -7,7 +7,8 @@
 import copy
 import time
 import uuid
-from concurrent.futures import ThreadPoolExecutor
+import datetime
+import concurrent.futures
 from azure.iot.hub import IoTHubRegistryManager
 from azure.iot.hub.protocol.models import Twin, TwinProperties, CloudToDeviceMethod
 from azure.eventhub import EventHubConsumerClient
@@ -78,8 +79,6 @@ def __init__(
         eventhub_connection_string,
         eventhub_consumer_group,
     ):
-        self._executor = ThreadPoolExecutor()
-
         self._registry_manager = IoTHubRegistryManager(iothub_connection_string)
 
         logger.info(
@@ -91,13 +90,32 @@ def __init__(
         self._eventhub_consumer_client = EventHubConsumerClient.from_connection_string(
             eventhub_connection_string, consumer_group=eventhub_consumer_group
         )
-
-        self._eventhub_future = self._executor.submit(self._eventhub_thread)
         self.device_id = None
         self.module_id = None
         self.incoming_patch_queue = queue.Queue()
         self.cv = threading.Condition()
         self.incoming_eventhub_events = {}
+        self._eventhub_start_position = datetime.datetime.now(datetime.timezone.utc)
+        self._eventhub_ready = threading.Event()
+        self._eventhub_future = self._start_eventhub_thread()
+        self._eventhub_future.add_done_callback(lambda _: self._eventhub_ready.set())
+
+    def _start_eventhub_thread(self):
+        future = concurrent.futures.Future()
+
+        def run():
+            if not future.set_running_or_notify_cancel():
+                return
+            try:
+                result = self._eventhub_thread()
+            except BaseException as e:
+                future.set_exception(e)
+            else:
+                future.set_result(result)
+
+        self._eventhub_worker = threading.Thread(target=run, daemon=True)
+        self._eventhub_worker.start()
+        return future
 
     def set_identity(self, device_id, module_id):
         if device_id != self.device_id or module_id != self.module_id:
@@ -113,6 +131,15 @@ def set_identity(self, device_id, module_id):
                     )
                 self.incoming_eventhub_events = {}
 
+    def wait_until_ready(self, timeout):
+        if not self._eventhub_ready.wait(timeout=timeout):
+            raise TimeoutError(
+                "EventHub consumer did not initialize within {} seconds".format(timeout)
+            )
+        if self._eventhub_future.done():
+            self._eventhub_future.result()
+            raise RuntimeError("EventHub consumer exited before initializing")
+
     def set_desired_properties(self, desired_props):
         if self.module_id:
             self._registry_manager.update_module_twin(
@@ -155,15 +182,26 @@ def send_c2d(self, payload, properties):
             raise TypeError("sending C2D to modules is not supported")
         self._registry_manager.send_c2d_message(self.device_id, payload, properties)
 
-    def wait_for_eventhub_arrival(self, message_id, timeout=900):
+    def wait_for_eventhub_arrival(self, message_id, timeout=900, event_filter=None):
         def get_event(inner_message_id):
             with self.cv:
                 arrivals = self.incoming_eventhub_events
 
                 # if message_id is not set, return any message
                 if not inner_message_id and len(arrivals):
-                    id = list(arrivals.keys())[0]
-                    logger.info("wait_for_eventhub_arrival(None) returning msgid={}".format(id))
+                    if event_filter:
+                        id = next(
+                            (
+                                event_id
+                                for event_id, event in arrivals.items()
+                                if event_filter(event)
+                            ),
+                            None,
+                        )
+                    else:
+                        id = next(iter(arrivals))
+                    if id:
+                        logger.info("wait_for_eventhub_arrival(None) returning msgid={}".format(id))
                 else:
                     id = inner_message_id
 
@@ -203,14 +241,13 @@ def shutdown(self):
 
         if self._eventhub_future:
             try:
-                # Ensure the EventHub receive loop exits before tearing down the executor
+                # Ensure the EventHub receive loop exits after closing the consumer
                 self._eventhub_future.result(timeout=30)
+            except concurrent.futures.TimeoutError:
+                logger.warning("_eventhub_thread did not exit within 30 seconds")
             except Exception:
                 logger.warning("_eventhub_thread did not exit cleanly", exc_info=True)
 
-        if self._executor:
-            self._executor.shutdown(wait=True)
-
     def _convert_incoming_event(self, event):
         try:
             event_body = event.body_as_json()
@@ -255,6 +292,8 @@ def on_partition_close(partition_context, reason):
 
         def on_event_batch(partition_context, events):
             try:
+                # A completed receive confirms that the Event Hub consumer is active
+                self._eventhub_ready.set()
                 for event in events:
                     device_id = get_device_id_from_event(event)
                     module_id = get_module_id_from_event(event)
@@ -294,6 +333,7 @@ def on_event_batch(partition_context, events):
                 logger.info("Starting EventHub receive")
                 self._eventhub_consumer_client.receive_batch(
                     max_wait_time=2,
+                    starting_position=self._eventhub_start_position,
                     on_event_batch=on_event_batch,
                     on_error=on_error,
                     on_partition_initialize=on_partition_initialize,

tests/e2e/iothub_e2e/aio/conftest.py

@@ -7,6 +7,7 @@
 import datetime
 import json
 import retry_async
+import const
 from utils import create_client_object
 from azure.iot.device.iothub.aio import IoTHubDeviceClient, IoTHubModuleClient
 
@@ -103,14 +104,16 @@ async def service_helper(executor):
         eventhub_consumer_group=test_env.EVENTHUB_CONSUMER_GROUP,
         executor=executor,
     )
-    yield service_helper
-
-    logger.info("----------------------------")
-    logger.info("shutting down service_helper")
-    logger.info("----------------------------")
+    try:
+        await service_helper.wait_until_ready(timeout=const.E2E_TIMEOUT)
+        yield service_helper
+    finally:
+        logger.info("----------------------------")
+        logger.info("shutting down service_helper")
+        logger.info("----------------------------")
 
-    await service_helper.shutdown()
+        await service_helper.shutdown()
 
-    logger.info("---------------------------------")
-    logger.info("service helper shut down complete")
-    logger.info("---------------------------------")
+        logger.info("---------------------------------")
+        logger.info("service helper shut down complete")
+        logger.info("---------------------------------")

tests/e2e/iothub_e2e/aio/test_c2d.py

@@ -6,6 +6,7 @@
 import logging
 import json
 from dev_utils import get_random_dict
+import const
 
 logger = logging.getLogger(__name__)
 logger.setLevel(level=logging.INFO)
@@ -38,7 +39,7 @@ async def handle_on_message_received(message):
 
         await service_helper.send_c2d(message, {})
 
-        await asyncio.wait_for(received.wait(), 60)
+        await asyncio.wait_for(received.wait(), timeout=const.E2E_TIMEOUT)
         assert received.is_set()
 
         assert received_message.data.decode("utf-8") == message

tests/e2e/iothub_e2e/aio/test_connect_disconnect.py

@@ -5,6 +5,8 @@
 import pytest
 import logging
 import parametrize
+import const
+import wait_helpers
 
 logger = logging.getLogger(__name__)
 logger.setLevel(level=logging.INFO)
@@ -78,7 +80,7 @@ async def handle_on_connection_state_change():
         await client.disconnect()
 
         # wait for handle_on_connection_state_change to reconnect
-        await reconnected_event.wait()
+        await asyncio.wait_for(reconnected_event.wait(), timeout=const.E2E_TIMEOUT)
 
         logger.info(
             "reconnect_event.wait() returned.  client.connected={}".format(client.connected)
@@ -160,7 +162,7 @@ async def handle_on_connection_state_change():
         await client.connect()
 
         # and wait for us to disconnect
-        await disconnected_event.wait()
+        await asyncio.wait_for(disconnected_event.wait(), timeout=const.E2E_TIMEOUT)
         assert not client.connected
 
         # sleep a while and make sure that we're still disconnected.
@@ -198,14 +200,16 @@ async def test_disconnect_on_drop_outgoing(self, client, dropper, leak_tracker):
         assert client.connected
         dropper.drop_outgoing()
 
-        while client.connected:
-            await asyncio.sleep(1)
+        await wait_helpers.async_wait_for_condition(
+            lambda: not client.connected, timeout=const.E2E_TIMEOUT
+        )
 
         # we've passed the test. Now wait to reconnect before we check for leaks. Otherwise we
         # have a pending ConnectOperation floating around and this would get tagged as a leak.
         dropper.restore_all()
-        while not client.connected:
-            await asyncio.sleep(1)
+        await wait_helpers.async_wait_for_condition(
+            lambda: client.connected, timeout=const.E2E_TIMEOUT
+        )
 
         leak_tracker.check_for_leaks()
 
@@ -221,13 +225,15 @@ async def test_disconnect_on_reject_outgoing(self, client, dropper, leak_tracker
         assert client.connected
         dropper.reject_outgoing()
 
-        while client.connected:
-            await asyncio.sleep(1)
+        await wait_helpers.async_wait_for_condition(
+            lambda: not client.connected, timeout=const.E2E_TIMEOUT
+        )
 
         # we've passed the test. Now wait to reconnect before we check for leaks. Otherwise we
         # have a pending ConnectOperation floating around and this would get tagged as a leak.
         dropper.restore_all()
-        while not client.connected:
-            await asyncio.sleep(1)
+        await wait_helpers.async_wait_for_condition(
+            lambda: client.connected, timeout=const.E2E_TIMEOUT
+        )
 
         leak_tracker.check_for_leaks()

tests/e2e/iothub_e2e/aio/test_methods.py

@@ -3,7 +3,7 @@
 # license information.
 import pytest
 import logging
-import asyncio
+
 import parametrize
 from dev_utils import get_random_dict
 from azure.iot.device.iothub import MethodResponse
@@ -62,7 +62,6 @@ async def handle_on_method_request_received(request):
             )
 
         client.on_method_request_received = handle_on_method_request_received
-        await asyncio.sleep(1)  # wait for subscribe, etc, to complete
 
         # invoke the method call
         method_response = await service_helper.invoke_method(method_name, request_payload)

tests/e2e/iothub_e2e/aio/test_sas_renewal.py

@@ -7,6 +7,7 @@
 import logging
 import test_config
 import parametrize
+import const
 
 logger = logging.getLogger(__name__)
 logger.setLevel(level=logging.INFO)
@@ -50,7 +51,7 @@ async def handle_on_connection_state_change():
         # setting on_connection_state_change seems to have the side effect of
         # calling handle_on_connection_state_change once with the initial value.
         # Wait for one disconnect/reconnect cycle so we can get past it.
-        await connected_event.wait()
+        await asyncio.wait_for(connected_event.wait(), timeout=const.E2E_TIMEOUT)
 
         # OK, we're ready to test.  wait for the renewal
         token_before_connect = str(token_object)
@@ -59,9 +60,9 @@ async def handle_on_connection_state_change():
         connected_event.clear()
 
         logger.info("Waiting for client to disconnect")
-        await disconnected_event.wait()
+        await asyncio.wait_for(disconnected_event.wait(), timeout=const.E2E_TIMEOUT)
         logger.info("Waiting for client to reconnect")
-        await connected_event.wait()
+        await asyncio.wait_for(connected_event.wait(), timeout=const.E2E_TIMEOUT)
         logger.info("Client reconnected")
 
         # Finally verify that our token changed.