mypy on tests

Author: stellasia

tests/unit/test_core_orchestrator.py

@@ -3,58 +3,53 @@
 
 
 @pytest.fixture(scope="function")
-def component():
-    return Component()
-
-
-@pytest.fixture(scope="function")
-def pipeline_branch(component):
+def pipeline_branch() -> Pipeline:
     pipe = Pipeline()
-    pipe.add_component("a", component)
-    pipe.add_component("b", component)
-    pipe.add_component("c", component)
+    pipe.add_component("a", Component())
+    pipe.add_component("b", Component())
+    pipe.add_component("c", Component())
     pipe.connect("a", "b")
     pipe.connect("a", "c")
     return pipe
 
 
 @pytest.fixture(scope="function")
-def pipeline_aggregation(component):
+def pipeline_aggregation() -> Pipeline:
     pipe = Pipeline()
-    pipe.add_component("a", component)
-    pipe.add_component("b", component)
-    pipe.add_component("c", component)
+    pipe.add_component("a", Component())
+    pipe.add_component("b", Component())
+    pipe.add_component("c", Component())
     pipe.connect("a", "b")
     pipe.connect("a", "c")
     return pipe
 
 
-def test_orchestrator_branch(pipeline_branch):
+async def test_orchestrator_branch(pipeline_branch: Pipeline) -> None:
     orchestrator = Orchestrator(pipeline=pipeline_branch)
     node_a = pipeline_branch.get_node_by_name("a")
-    node_a.status = RunStatus.DONE
-    next_tasks = orchestrator.next(node_a)
+    node_a.status = RunStatus.DONE  # type: ignore
+    next_tasks = [n async for n in orchestrator.next(node_a)]  # type: ignore
     next_task_names = [n.name for n in next_tasks]
     assert next_task_names == ["b", "c"]
 
 
-def test_orchestrator_aggregation(pipeline_aggregation):
+async def test_orchestrator_aggregation(pipeline_aggregation: Pipeline) -> None:
     orchestrator = Orchestrator(pipeline=pipeline_aggregation)
     node_a = pipeline_aggregation.get_node_by_name("a")
-    node_a.status = RunStatus.DONE
+    node_a.status = RunStatus.DONE  # type: ignore
     node_b = pipeline_aggregation.get_node_by_name("b")
-    node_b.status = RunStatus.DONE
-    next_tasks = orchestrator.next(node_a)
+    node_b.status = RunStatus.DONE  # type: ignore
+    next_tasks = [n async for n in orchestrator.next(node_a)]  # type: ignore
     next_task_names = [n.name for n in next_tasks]
     assert next_task_names == ["c"]
 
 
-def test_orchestrator_aggregation_waiting(pipeline_aggregation):
+async def test_orchestrator_aggregation_waiting(pipeline_aggregation: Pipeline) -> None:
     orchestrator = Orchestrator(pipeline=pipeline_aggregation)
     node_a = pipeline_aggregation.get_node_by_name("a")
-    node_a.status = RunStatus.DONE
+    node_a.status = RunStatus.DONE  # type: ignore
     node_b = pipeline_aggregation.get_node_by_name("a")
-    node_b.status = RunStatus.UNKNOWN
-    next_tasks = orchestrator.next(node_a)
+    node_b.status = RunStatus.UNKNOWN  # type: ignore
+    next_tasks = [n async for n in orchestrator.next(node_a)]  # type: ignore
     next_task_names = [n.name for n in next_tasks]
     assert next_task_names == []

tests/unit/test_core_pipeline.py

@@ -6,29 +6,8 @@
 from neo4j_genai.core.pipeline import Component, Pipeline
 
 
-@pytest.fixture(scope="function")
-def component_multiply():
-    class ComponentMultiply(Component):
-        def __init__(self, r: float = 2.0) -> None:
-            self.r = r
-
-        async def run(self, number: float):
-            return {"product": number * self.r}
-
-    return ComponentMultiply()
-
-
-@pytest.fixture(scope="function")
-def component_add():
-    class ComponentAdd(Component):
-        async def run(self, number1: float, number2: float):
-            return {"sum": number1 + number2}
-
-    return ComponentAdd()
-
-
 @pytest.mark.asyncio
-async def test_simple_pipeline_two_components():
+async def test_simple_pipeline_two_components() -> None:
     pipe = Pipeline()
     component_a = AsyncMock(spec=Component)
     component_a.run = AsyncMock(return_value={})
@@ -45,7 +24,7 @@ async def test_simple_pipeline_two_components():
 
 
 @pytest.mark.asyncio
-async def test_simple_pipeline_two_components_parameter_propagation():
+async def test_simple_pipeline_two_components_parameter_propagation() -> None:
     pipe = Pipeline()
     component_a = AsyncMock(spec=Component)
     component_a.run = AsyncMock(return_value={"product": 20})
@@ -68,7 +47,7 @@ async def test_simple_pipeline_two_components_parameter_propagation():
 
 
 @pytest.mark.asyncio
-async def test_pipeline_branches():
+async def test_pipeline_branches() -> None:
     pipe = Pipeline()
     component_a = AsyncMock(spec=Component)
     component_a.run = AsyncMock(return_value={})
@@ -88,7 +67,7 @@ async def test_pipeline_branches():
 
 
 @pytest.mark.asyncio
-async def test_pipeline_aggregation():
+async def test_pipeline_aggregation() -> None:
     pipe = Pipeline()
     component_a = AsyncMock(spec=Component)
     component_a.run = AsyncMock(return_value={})