test: Expand cancellation unit test to both NATS and TCP request plane

Author: kthui

lib/bindings/python/tests/cancellation/test_cancellation.py

@@ -165,6 +165,7 @@ async def client(runtime, namespace):
 
 @pytest.mark.forked
 @pytest.mark.asyncio
+@pytest.mark.parametrize("request_plane", ["nats", "tcp"], indirect=True)
 async def test_client_context_cancel(temp_file_store, server, client):
     _, handler = server
     context = Context()
@@ -198,6 +199,7 @@ async def test_client_context_cancel(temp_file_store, server, client):
 
 @pytest.mark.forked
 @pytest.mark.asyncio
+@pytest.mark.parametrize("request_plane", ["nats", "tcp"], indirect=True)
 async def test_client_loop_break(temp_file_store, server, client):
     _, handler = server
     stream = await client.generate("_generate_until_context_cancelled")
@@ -230,6 +232,7 @@ async def test_client_loop_break(temp_file_store, server, client):
 
 @pytest.mark.forked
 @pytest.mark.asyncio
+@pytest.mark.parametrize("request_plane", ["nats", "tcp"], indirect=True)
 async def test_server_context_cancel(temp_file_store, server, client):
     _, handler = server
     stream = await client.generate("_generate_and_cancel_context")
@@ -254,6 +257,7 @@ async def test_server_context_cancel(temp_file_store, server, client):
 
 @pytest.mark.forked
 @pytest.mark.asyncio
+@pytest.mark.parametrize("request_plane", ["nats", "tcp"], indirect=True)
 async def test_server_raise_cancelled(temp_file_store, server, client):
     _, handler = server
     stream = await client.generate("_generate_and_raise_cancelled")
@@ -282,6 +286,7 @@ async def test_server_raise_cancelled(temp_file_store, server, client):
 
 @pytest.mark.forked
 @pytest.mark.asyncio
+@pytest.mark.parametrize("request_plane", ["nats", "tcp"], indirect=True)
 async def test_client_context_already_cancelled(temp_file_store, server, client):
     _, handler = server
     context = Context()
@@ -304,6 +309,7 @@ async def test_client_context_already_cancelled(temp_file_store, server, client)
 
 @pytest.mark.forked
 @pytest.mark.asyncio
+@pytest.mark.parametrize("request_plane", ["nats", "tcp"], indirect=True)
 async def test_client_context_cancel_before_await_request(
     temp_file_store, server, client
 ):

lib/bindings/python/tests/conftest.py

@@ -402,8 +402,34 @@ def temp_file_store():
         yield tmpdir
 
 
+@pytest.fixture
+def store_kv(request):
+    """
+    KV store for runtime. Defaults to "file".
+
+    To iterate over multiple stores in a test:
+        @pytest.mark.parametrize("store_kv", ["file", "etcd"], indirect=True)
+        async def test_example(runtime):
+            ...
+    """
+    return getattr(request, "param", "file")
+
+
+@pytest.fixture
+def request_plane(request):
+    """
+    Request plane for runtime. Defaults to "nats".
+
+    To iterate over multiple transports in a test:
+        @pytest.mark.parametrize("request_plane", ["tcp", "nats"], indirect=True)
+        async def test_example(runtime):
+            ...
+    """
+    return getattr(request, "param", "nats")
+
+
 @pytest.fixture(scope="function", autouse=False)
-async def runtime(request):
+async def runtime(request, store_kv, request_plane):
     """
     Create a DistributedRuntime for testing.
 
@@ -413,6 +439,14 @@ async def runtime(request):
 
     Without @pytest.mark.forked in isolated mode, you will get "Worker already initialized"
     errors when multiple tests try to create runtimes in the same process.
+
+    The store_kv and request_plane can be customized by overriding their fixtures
+    or using @pytest.mark.parametrize with indirect=True:
+
+        @pytest.mark.forked
+        @pytest.mark.parametrize("store_kv", ["etcd"], indirect=True)
+        async def test_with_etcd(runtime):
+            ...
     """
     # Check if the test is marked with @pytest.mark.forked (only in isolated mode)
     if ENABLE_ISOLATED_ETCD_AND_NATS:
@@ -435,6 +469,6 @@ async def test_my_test(runtime):
             )
 
     loop = asyncio.get_running_loop()
-    runtime = DistributedRuntime(loop, "file", "nats")
+    runtime = DistributedRuntime(loop, store_kv, request_plane)
     yield runtime
     runtime.shutdown()