Parametrize with async for style tests

tests/test_2way.py

@@ -1,7 +1,8 @@
 """
 Bidirectional streaming and context API.
-"""
 
+"""
+import pytest
 import trio
 import tractor
 
@@ -86,7 +87,9 @@ async def simple_rpc(
     data: int,
 
 ) -> None:
+    """Test a small ping-pong server.
 
+    """
     # signal to parent that we're up
     await ctx.started(data + 1)
 
@@ -106,7 +109,44 @@ async def simple_rpc(
                 count += 1
 
 
-def test_simple_rpc():
+@tractor.context
+async def simple_rpc_with_forloop(
+
+    ctx: tractor.Context,
+    data: int,
+
+) -> None:
+    """Same as previous test but using ``async for`` syntax/api.
+
+    """
+
+    # signal to parent that we're up
+    await ctx.started(data + 1)
+
+    print('opening stream in callee')
+    async with ctx.open_stream() as stream:
+
+        count = 0
+        async for msg in stream:
+
+            assert msg == 'ping'
+            print('pong')
+            await stream.send('pong')
+            count += 1
+
+        else:
+            assert count == 10
+
+
+@pytest.mark.parametrize(
+    'use_async_for',
+    [True, False],
+)
+@pytest.mark.parametrize(
+    'server_func',
+    [simple_rpc, simple_rpc_with_forloop],
+)
+def test_simple_rpc(server_func, use_async_for):
     """The simplest request response pattern.
 
     """
@@ -119,7 +159,7 @@ def test_simple_rpc():
             )
 
             async with portal.open_context(
-                simple_rpc,
+                server_func,  # taken from pytest parameterization
                 data=10,
             ) as (ctx, sent):
 
@@ -127,11 +167,29 @@ def test_simple_rpc():
 
                 async with ctx.open_stream() as stream:
 
-                    for _ in range(10):
+                    if use_async_for:
 
+                        count = 0
+                        # receive msgs using async for style
                         print('ping')
                         await stream.send('ping')
-                        assert await stream.receive() == 'pong'
+
+                        async for msg in stream:
+                            assert msg == 'pong'
+                            print('ping')
+                            await stream.send('ping')
+                            count += 1
+
+                            if count >= 9:
+                                break
+
+                    else:
+                        # classic send/receive style
+                        for _ in range(10):
+
+                            print('ping')
+                            await stream.send('ping')
+                            assert await stream.receive() == 'pong'
 
                 # stream should terminate here