diff --git a/bellows/ezsp/__init__.py b/bellows/ezsp/__init__.py index 29554c93..b55a71bb 100644 --- a/bellows/ezsp/__init__.py +++ b/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: diff --git a/bellows/thread.py b/bellows/thread.py index 4311768d..5789c2d0 100644 --- a/bellows/thread.py +++ b/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 diff --git a/tests/test_ezsp.py b/tests/test_ezsp.py index d548309a..17e16d59 100644 --- a/tests/test_ezsp.py +++ b/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) diff --git a/tests/test_thread.py b/tests/test_thread.py index 72efa701..db92e69a 100644 --- a/tests/test_thread.py +++ b/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()