aio.run.runner: Fix for cleanup decorator (#2336)

Signed-off-by: Ryan Northey <[email protected]>

aio.run.runner/VERSION

@@ -1 +1 @@
-0.3.4-dev
+0.3.4

aio.run.runner/aio/run/runner/decorators.py

@@ -50,11 +50,11 @@ def cleansup(fun) -> Callable:
     """
 
     @wraps(fun)
-    def wrapped(self, *args, **kwargs) -> int | None:
+    async def wrapped(self, *args, **kwargs) -> int | None:
         try:
-            return fun(self, *args, **kwargs)
+            return await fun(self, *args, **kwargs)
         finally:
-            self.cleanup()
+            await self.cleanup()
 
     # mypy doesnt trust `@wraps` to give back a `__wrapped__` object so we
     # need to code defensively here

aio.run.runner/tests/test_decorators.py

@@ -1,5 +1,5 @@
 
-from unittest.mock import MagicMock, PropertyMock
+from unittest.mock import AsyncMock, MagicMock, PropertyMock
 
 import pytest
 
@@ -125,7 +125,7 @@ async def test_cleansup(iters, raises):
 
     assert run.run.__wrapped__.__cleansup__ is True
 
-    run.cleanup = MagicMock()
+    run.cleanup = AsyncMock()
     if raises:
         with pytest.raises(Exception):
             await run.run(*args, **kwargs)