Add tests

Author: stellasia

src/neo4j_graphrag/experimental/pipeline/component.py

@@ -81,7 +81,7 @@ class Component(abc.ABC, metaclass=ComponentMeta):
     # added here for the type checker
     # DO NOT CHANGE
     component_inputs: dict[str, dict[str, str | bool]]
-    component_outputs: dict[str, dict[str, str | bool]]
+    component_outputs: dict[str, dict[str, str | bool | type]]
 
     @abc.abstractmethod
     async def run(self, *args: Any, **kwargs: Any) -> DataModel:

tests/unit/experimental/pipeline/components.py

@@ -13,6 +13,7 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 from neo4j_graphrag.experimental.pipeline import Component, DataModel
+from neo4j_graphrag.experimental.pipeline.types.context import RunContext
 
 
 class StringResultModel(DataModel):
@@ -41,3 +42,13 @@ async def run(self, number1: int, number2: int) -> IntResultModel:
 class ComponentMultiply(Component):
     async def run(self, number1: int, number2: int = 2) -> IntResultModel:
         return IntResultModel(result=number1 * number2)
+
+
+class ComponentMultiplyWithContext(Component):
+    async def run_with_context(
+        self, context_: RunContext, number1: int, number2: int = 2
+    ) -> IntResultModel:
+        await context_.notify(
+            message="my message", data={"number1": number1, "number2": number2}
+        )
+        return IntResultModel(result=number1 * number2)

tests/unit/experimental/pipeline/test_component.py

@@ -12,7 +12,12 @@
 #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
-from .components import ComponentMultiply
+from unittest.mock import MagicMock, AsyncMock
+
+import pytest
+
+from neo4j_graphrag.experimental.pipeline.types.context import RunContext
+from .components import ComponentMultiply, ComponentMultiplyWithContext, IntResultModel
 
 
 def test_component_inputs() -> None:
@@ -26,3 +31,41 @@ def test_component_inputs() -> None:
 def test_component_outputs() -> None:
     outputs = ComponentMultiply.component_outputs
     assert "result" in outputs
+    assert outputs["result"]["has_default"] is True
+    assert outputs["result"]["annotation"] == int
+
+
+@pytest.mark.asyncio
+async def test_component_run() -> None:
+    c = ComponentMultiply()
+    result = await c.run(number1=1, number2=2)
+    assert isinstance(result, IntResultModel)
+    assert isinstance(
+        result.result, ComponentMultiply.component_outputs["result"]["annotation"]
+    )
+
+
+@pytest.mark.asyncio
+async def test_component_run_with_context_default_implementation() -> None:
+    c = ComponentMultiply()
+    result = await c.run_with_context(
+        # context can not be null in the function signature,
+        # but it's ignored in this case
+        None,  # type: ignore
+        number1=1,
+        number2=2,
+    )
+    assert result.result == 2
+
+
+@pytest.mark.asyncio
+async def test_component_run_with_context() -> None:
+    c = ComponentMultiplyWithContext()
+    notifier_mock = AsyncMock()
+    result = await c.run_with_context(
+        RunContext(run_id="run_id", task_name="task_name", notifier=notifier_mock),
+        number1=1,
+        number2=2,
+    )
+    assert result.result == 2
+    notifier_mock.assert_awaited_once()