[TP-1786] Create subtests for each channel in client tests

Author: mimarcel

tests/common.py

@@ -1,4 +1,5 @@
 import os
+import pytest
 import time
 from datetime import datetime
 from typing import List, Tuple
@@ -65,41 +66,58 @@ def metadata(pat: bool = False) -> Tuple[Tuple[str, str], Tuple[str, str]]:
     )
 
 
-def grpc_channel(func):
+def run_tests_using_channels(func, channel_keys):
     """
-    A decorator that runs the test using the gRPC channel.
+    A decorator that runs the test using given channels.
     :param func: The test function.
+    :param channel_keys: A list of channel keys to use for the test. A subtest is created for each channel.
     :return: A function wrapper.
     """
 
-    def func_wrapper():
+    @pytest.mark.parametrize('channel_key', channel_keys)
+    def func_wrapper(channel_key):
+        return func(get_channel(channel_key))
+
+    return func_wrapper
+
+
+def get_channel(channel_key):
+    if channel_key == "grpc":
         if os.getenv("CLARIFAI_GRPC_INSECURE", "False").lower() in ("true", "1", "t"):
-            channel = ClarifaiChannel.get_insecure_grpc_channel(port=443)
+            return ClarifaiChannel.get_insecure_grpc_channel(port=443)
         else:
-            channel = ClarifaiChannel.get_grpc_channel()
-        func(channel)
+            return ClarifaiChannel.get_grpc_channel()
 
-    return func_wrapper
+    if channel_key == "json":
+        return ClarifaiChannel.get_json_channel()
 
+    if channel_key == "asyncio":
+        if os.getenv("CLARIFAI_GRPC_INSECURE", "False").lower() in ("true", "1", "t"):
+            return ClarifaiChannel.get_aio_insecure_grpc_channel(port=443)
+        else:
+            return ClarifaiChannel.get_aio_grpc_channel()
 
-def both_channels(func):
+    raise ValueError(f"Unknown channel {channel_key}")
+
+
+def grpc_channel(func):
     """
-    A decorator that runs the test first using the gRPC channel and then using the JSON channel.
+    A decorator that runs the test using the gRPC channel.
     :param func: The test function.
     :return: A function wrapper.
     """
 
-    def func_wrapper():
-        if os.getenv("CLARIFAI_GRPC_INSECURE", "False").lower() in ("true", "1", "t"):
-            channel = ClarifaiChannel.get_insecure_grpc_channel(port=443)
-        else:
-            channel = ClarifaiChannel.get_grpc_channel()
-        func(channel)
+    return run_tests_using_channels(func, ["grpc"])
 
-        channel = ClarifaiChannel.get_json_channel()
-        func(channel)
 
-    return func_wrapper
+def both_channels(func):
+    """
+    A decorator that runs the test first using the gRPC channel and then using the JSON channel.
+    :param func: The test function.
+    :return: A function wrapper.
+    """
+
+    return run_tests_using_channels(func, ["grpc", "json"])
 
 
 def asyncio_channel(func):
@@ -109,13 +127,7 @@ def asyncio_channel(func):
     :return: A function wrapper.
     """
 
-    async def func_wrapper():
-        channel = ClarifaiChannel.get_aio_grpc_channel()
-        if os.getenv("CLARIFAI_GRPC_INSECURE", "False").lower() in ("true", "1", "t"):
-            channel = ClarifaiChannel.get_aio_insecure_grpc_channel(port=443)
-        await func(channel)
-
-    return func_wrapper
+    return run_tests_using_channels(func, ["asyncio"])
 
 
 def wait_for_inputs_upload(stub, metadata, input_ids):