Add unit tests

Author: gmazoyer

infrahub_sdk/batch.py

@@ -27,9 +27,9 @@ def execute(self, return_exceptions: bool = False) -> tuple[Optional[InfrahubNod
         result = None
         try:
             result = self.task(*self.args, **self.kwargs)
-        except Exception as exc:
+        except Exception as exc:  # pylint: disable=broad-exception-caught
             if return_exceptions:
-                return (self.node, exc)
+                return self.node, exc
             raise exc
 
         return self.node, result
@@ -101,7 +101,7 @@ def add(self, *args: Any, task: Callable[..., Any], node: Optional[Any] = None,
 
     def execute(self) -> Generator[tuple[Optional[InfrahubNodeSync], Any], None, None]:
         with ThreadPoolExecutor(max_workers=self.max_concurrent_execution) as executor:
-            futures = [executor.submit(task.execute) for task in self._tasks]
+            futures = [executor.submit(task.execute, return_exceptions=self.return_exceptions) for task in self._tasks]
             for future in futures:
                 node, result = future.result()
                 if isinstance(result, Exception) and not self.return_exceptions:

tests/unit/sdk/test_batch.py

@@ -11,43 +11,123 @@
 
     from tests.unit.sdk.conftest import BothClients
 
+client_types = ["standard", "sync"]
 
+
+@pytest.mark.parametrize("client_type", client_types)
+async def test_batch_execution(clients: BothClients, client_type: str):
+    r: list[int] = []
+    tasks_number = 10
+
+    if client_type == "standard":
+
+        async def test_func() -> int:
+            return 1
+
+        batch = await clients.standard.create_batch()
+        for _ in range(tasks_number):
+            batch.add(task=test_func)
+
+        assert batch.num_tasks == tasks_number
+        async for _, result in batch.execute():
+            r.append(result)
+    else:
+
+        def test_func() -> int:
+            return 1
+
+        batch = clients.sync.create_batch()
+        for _ in range(tasks_number):
+            batch.add(task=test_func)
+
+        assert batch.num_tasks == tasks_number
+        for _, result in batch.execute():
+            r.append(result)
+
+    assert r == [1] * tasks_number
+
+
+@pytest.mark.parametrize("client_type", client_types)
 async def test_batch_return_exception(
-    httpx_mock: HTTPXMock, mock_query_mutation_location_create_failed, mock_schema_query_01, clients: BothClients
+    httpx_mock: HTTPXMock,
+    mock_query_mutation_location_create_failed,
+    mock_schema_query_01,
+    clients: BothClients,
+    client_type: str,
 ):  # pylint: disable=unused-argument
-    batch = await clients.standard.create_batch(return_exceptions=True)
-    locations = ["JFK1", "JFK1"]
-    results = []
-    for location_name in locations:
-        data = {"name": {"value": location_name, "is_protected": True}}
-        obj = await clients.standard.create(kind="BuiltinLocation", data=data)
-        batch.add(task=obj.save, node=obj)
-        results.append(obj)
-
-    result_iter = batch.execute()
-    # Assert first node success
-    node, result = await result_iter.__anext__()
-    assert node == results[0]
-    assert not isinstance(result, Exception)
-
-    # Assert second node failure
-    node, result = await result_iter.__anext__()
-    assert node == results[1]
-    assert isinstance(result, GraphQLError)
-    assert "An error occurred while executing the GraphQL Query" in str(result)
+    if client_type == "standard":
+        batch = await clients.standard.create_batch(return_exceptions=True)
+        locations = ["JFK1", "JFK1"]
+        results = []
+        for location_name in locations:
+            data = {"name": {"value": location_name, "is_protected": True}}
+            obj = await clients.standard.create(kind="BuiltinLocation", data=data)
+            batch.add(task=obj.save, node=obj)
+            results.append(obj)
+
+        result_iter = batch.execute()
+        # Assert first node success
+        node, result = await result_iter.__anext__()
+        assert node == results[0]
+        assert not isinstance(result, Exception)
+
+        # Assert second node failure
+        node, result = await result_iter.__anext__()
+        assert node == results[1]
+        assert isinstance(result, GraphQLError)
+        assert "An error occurred while executing the GraphQL Query" in str(result)
+    else:
+        batch = clients.sync.create_batch(return_exceptions=True)
+        locations = ["JFK2", "JFK2"]
+        results = []
+        for location_name in locations:
+            data = {"name": {"value": location_name, "is_protected": True}}
+            obj = clients.sync.create(kind="BuiltinLocation", data=data)
+            batch.add(task=obj.save, node=obj)
+            results.append(obj)
 
+        result_iter = batch.execute()
+        # Assert first node success
+        node, result = next(result_iter)
+        assert node == results[0]
+        assert not isinstance(result, Exception)
 
+        # Assert second node failure
+        node, result = next(result_iter)
+        assert node == results[1]
+        assert isinstance(result, GraphQLError)
+        assert "An error occurred while executing the GraphQL Query" in str(result)
+
+
+@pytest.mark.parametrize("client_type", client_types)
 async def test_batch_exception(
-    httpx_mock: HTTPXMock, mock_query_mutation_location_create_failed, mock_schema_query_01, clients: BothClients
+    httpx_mock: HTTPXMock,
+    mock_query_mutation_location_create_failed,
+    mock_schema_query_01,
+    clients: BothClients,
+    client_type: str,
 ):  # pylint: disable=unused-argument
-    batch = await clients.standard.create_batch(return_exceptions=False)
-    locations = ["JFK1", "JFK1"]
-    for location_name in locations:
-        data = {"name": {"value": location_name, "is_protected": True}}
-        obj = await clients.standard.create(kind="BuiltinLocation", data=data)
-        batch.add(task=obj.save, node=obj)
-
-    with pytest.raises(GraphQLError) as exc:
-        async for _, _ in batch.execute():
-            pass
-    assert "An error occurred while executing the GraphQL Query" in str(exc.value)
+    if client_type == "standard":
+        batch = await clients.standard.create_batch(return_exceptions=False)
+        locations = ["JFK1", "JFK1"]
+        for location_name in locations:
+            data = {"name": {"value": location_name, "is_protected": True}}
+            obj = await clients.standard.create(kind="BuiltinLocation", data=data)
+            batch.add(task=obj.save, node=obj)
+
+        with pytest.raises(GraphQLError) as exc:
+            async for _, _ in batch.execute():
+                pass
+        assert "An error occurred while executing the GraphQL Query" in str(exc.value)
+    else:
+        batch = clients.sync.create_batch(return_exceptions=False)
+        locations = ["JFK2", "JFK2"]
+        for location_name in locations:
+            data = {"name": {"value": location_name, "is_protected": True}}
+            obj = clients.sync.create(kind="BuiltinLocation", data=data)
+            batch.add(task=obj.save, node=obj)
+
+        with pytest.raises(GraphQLError) as exc:
+            for _, _ in batch.execute():
+                pass
+        assert "An error occurred while executing the GraphQL Query" in str(exc.value)