service helper changes test

Author: cartertinney

dev_utils/dev_utils/service_helper.py

@@ -15,21 +15,29 @@ def __init__(
         event_loop=None,
         executor=None,
     ):
-        self._event_loop = event_loop or asyncio.get_event_loop()
+        # Keep the passed-in loop as a fallback, but prefer the currently running loop when
+        # scheduling work so we never attach futures to a different loop than the caller.
+        self._event_loop_override = event_loop
         self._executor = executor or concurrent.futures.ThreadPoolExecutor()
         self._inner_object = ServiceHelperSync(
             iothub_connection_string, eventhub_connection_string, eventhub_consumer_group
         )
 
+    def _get_event_loop(self):
+        try:
+            return asyncio.get_running_loop()
+        except RuntimeError:
+            return self._event_loop_override or asyncio.get_event_loop()
+
+    async def _run_in_executor(self, func, *args):
+        loop = self._get_event_loop()
+        return await loop.run_in_executor(self._executor, func, *args)
+
     def set_identity(self, device_id, module_id):
         return self._inner_object.set_identity(device_id, module_id)
 
     async def set_desired_properties(self, desired_props):
-        return await self._event_loop.run_in_executor(
-            self._executor,
-            self._inner_object.set_desired_properties,
-            desired_props,
-        )
+        return await self._run_in_executor(self._inner_object.set_desired_properties, desired_props)
 
     async def invoke_method(
         self,
@@ -38,8 +46,7 @@ async def invoke_method(
         connect_timeout_in_seconds=None,
         response_timeout_in_seconds=None,
     ):
-        return await self._event_loop.run_in_executor(
-            self._executor,
+        return await self._run_in_executor(
             self._inner_object.invoke_method,
             method_name,
             payload,
@@ -52,25 +59,21 @@ async def send_c2d(
         payload,
         properties,
     ):
-        return await self._event_loop.run_in_executor(
-            self._executor, self._inner_object.send_c2d, payload, properties
-        )
+        return await self._run_in_executor(self._inner_object.send_c2d, payload, properties)
 
     async def wait_for_eventhub_arrival(self, message_id, timeout=60):
-        return await self._event_loop.run_in_executor(
-            self._executor,
+        return await self._run_in_executor(
             self._inner_object.wait_for_eventhub_arrival,
             message_id,
             timeout,
         )
 
     async def get_next_reported_patch_arrival(self, block=True, timeout=240):
-        return await self._event_loop.run_in_executor(
-            self._executor,
+        return await self._run_in_executor(
             self._inner_object.get_next_reported_patch_arrival,
             block,
             timeout,
         )
 
     async def shutdown(self):
-        return await self._event_loop.run_in_executor(self._executor, self._inner_object.shutdown)
+        return await self._run_in_executor(self._inner_object.shutdown)