Async test helpers now preserve type info

Author: dpwatrous

sdk/batch/azure-batch/tests/async_wrapper.py

@@ -1,19 +1,22 @@
 from collections.abc import AsyncIterable
 import inspect
+from typing import Any, cast, Coroutine, Iterable, TypeVar, Union
 
+T = TypeVar("T")
 
-# wrapper to handle async and sync objects
-async def async_wrapper(obj):
+async def wrap_result(result: Union[T, Coroutine[Any, Any, T]]) -> T:
+    """Handle an non-list operation result and await it if it's a coroutine"""
+    if inspect.iscoroutine(result):
+        result = await result
+        return await wrap_result(result)
+    return cast(T, result)
 
-    if isinstance(obj, AsyncIterable):
-        items = []
-        async for item in obj:
+async def wrap_list_result(result: Union[Iterable[T], AsyncIterable[T]]) -> Iterable[T]:
+    """Handle a list operation result and convert to a list if it's an AsyncIterable"""    
+    if isinstance(result, AsyncIterable):
+        items: Iterable[T] = []
+        async for item in result:
             items.append(item)
         return items
-
-    if inspect.iscoroutine(obj):
-        waited = await obj
-        # wrap again to handle nested coroutines or async generators
-        return await async_wrapper(waited)
-
-    return obj
+    
+    return result

sdk/batch/azure-batch/tests/decorators.py

@@ -19,7 +19,7 @@
     stop_record_or_playback,
 )
 
-from async_wrapper import async_wrapper
+from async_wrapper import wrap_result
 
 
 # A modified version of devtools_testutils.aio.recorded_by_proxy_async
@@ -149,6 +149,6 @@ async def wrapper(self, BatchClient, **kwargs):
         except Exception as err:
             raise err
         finally:
-            await async_wrapper(client.close())
+            await wrap_result(client.close())
 
     return wrapper

sdk/batch/azure-batch/tests/test_async_wrapper.py

@@ -1,4 +1,4 @@
-from async_wrapper import async_wrapper
+from async_wrapper import wrap_list_result, wrap_result
 import pytest
 
 
@@ -8,7 +8,7 @@ async def test_iscoroutine(self):
         async def func():
             return 1
 
-        result = await async_wrapper(func())
+        result = await wrap_result(func())
         assert result == 1
 
     @pytest.mark.asyncio
@@ -17,19 +17,7 @@ async def func():
             for i in range(3):
                 yield i
 
-        result = await async_wrapper(func())
-        assert result == [0, 1, 2]
-
-    @pytest.mark.asyncio
-    async def test_isNestedAsyncIterable(self):
-        async def func():
-            async def nested():
-                for i in range(3):
-                    yield i
-
-            return nested()
-
-        result = await async_wrapper(func())
+        result = await wrap_list_result(func())
         assert result == [0, 1, 2]
 
     @pytest.mark.asyncio
@@ -40,7 +28,7 @@ async def nested():
 
             return nested()
 
-        result = await async_wrapper(func())
+        result = await wrap_result(func())
         assert result == 2
 
     @pytest.mark.asyncio
@@ -50,12 +38,12 @@ def func():
                 yield i
 
         iterable = func()
-        result = await async_wrapper(iterable)
+        result = await wrap_list_result(iterable)
         assert result == iterable
 
     @pytest.mark.asyncio
     async def test_isSync(self):
         def func():
             return 1
 
-        assert await async_wrapper(func()) == 1
+        assert await wrap_result(func()) == 1