Handle async proxy calls after event loop closes

bellows/ezsp/__init__.py

@@ -221,7 +221,11 @@ async def get_xncp_features(self) -> xncp.FirmwareFeatures:
     async def disconnect(self):
         self.stop_ezsp()
         if self._gw:
-            await self._gw.disconnect()
+            try:
+                await self._gw.disconnect()
+            except ConnectionResetError:
+                # A closed gateway loop is already disconnected.
+                pass
             self._gw = None
 
     async def _command(self, name: str, *args: Any, **kwargs: Any) -> Any:

bellows/thread.py

@@ -1,6 +1,7 @@
 import asyncio
 from concurrent.futures import ThreadPoolExecutor
 import functools
+import inspect
 import logging
 
 LOGGER = logging.getLogger(__name__)
@@ -88,6 +89,23 @@ def __getattr__(self, name):
                 )
             )
 
+        if inspect.iscoroutinefunction(func):
+
+            async def async_func_wrapper(*args, **kwargs):
+                """Run async proxy calls when this returned coroutine is awaited."""
+                loop = self._obj_loop
+                curr_loop = asyncio.get_running_loop()
+                call = functools.partial(func, *args, **kwargs)
+                if loop == curr_loop:
+                    return await call()
+                if loop.is_closed():
+                    LOGGER.warning("Attempted to use a closed event loop")
+                    raise ConnectionResetError("Attempted to use a closed event loop")
+                future = asyncio.run_coroutine_threadsafe(call(), loop)
+                return await asyncio.wrap_future(future, loop=curr_loop)
+
+            return async_func_wrapper
+
         def func_wrapper(*args, **kwargs):
             loop = self._obj_loop
             curr_loop = asyncio.get_running_loop()
@@ -98,21 +116,17 @@ def func_wrapper(*args, **kwargs):
                 # Disconnected
                 LOGGER.warning("Attempted to use a closed event loop")
                 return
-            if asyncio.iscoroutinefunction(func):
-                future = asyncio.run_coroutine_threadsafe(call(), loop)
-                return asyncio.wrap_future(future, loop=curr_loop)
-            else:
-
-                def check_result_wrapper():
-                    result = call()
-                    if result is not None:
-                        raise TypeError(
-                            (
-                                "ThreadsafeProxy can only wrap functions with no return"
-                                "value \nUse an async method to return values: {}.{}"
-                            ).format(self._obj.__class__.__name__, name)
-                        )
-
-                loop.call_soon_threadsafe(check_result_wrapper)
+
+            def check_result_wrapper():
+                result = call()
+                if result is not None:
+                    raise TypeError(
+                        (
+                            "ThreadsafeProxy can only wrap functions with no return"
+                            "value \nUse an async method to return values: {}.{}"
+                        ).format(self._obj.__class__.__name__, name)
+                    )
+
+            loop.call_soon_threadsafe(check_result_wrapper)
 
         return func_wrapper

tests/test_ezsp.py

@@ -81,6 +81,14 @@ async def test_disconnect(ezsp_f):
     assert len(gw_disconnect.mock_calls) == 1
 
 
+async def test_disconnect_closed_gateway_loop(ezsp_f):
+    ezsp_f._gw.disconnect = AsyncMock(side_effect=ConnectionResetError)
+
+    await ezsp_f.disconnect()
+
+    assert ezsp_f._gw is None
+
+
 def test_attr(ezsp_f):
     m = ezsp_f.getValue
     assert isinstance(m, functools.partial)

tests/test_thread.py

@@ -119,7 +119,9 @@ async def magic():
         return mock.sentinel.result
 
     obj.test = magic
-    result = await proxy.test()
+    call = proxy.test()
+    assert call_count == 0
+    result = await call
 
     assert call_count == 1
     assert result == mock.sentinel.result
@@ -161,6 +163,20 @@ async def test_proxy_loop_closed():
     assert obj.test.call_count == 0
 
 
+async def test_proxy_async_loop_closed():
+    loop = asyncio.new_event_loop()
+    obj = mock.MagicMock()
+
+    obj.test = mock.AsyncMock()
+    proxy = ThreadsafeProxy(obj, loop)
+    loop.close()
+
+    with pytest.raises(ConnectionResetError, match="closed event loop"):
+        await proxy.test()
+
+    obj.test.assert_not_awaited()
+
+
 async def test_thread_task_cancellation_after_stop(thread):
     loop = asyncio.get_event_loop()
     obj = mock.MagicMock()