참고소스 수정본

참고/neo4j-graphrag-python-main/tests/unit/experimental/pipeline/test_component.py

@@ -0,0 +1,89 @@
+#  Copyright (c) "Neo4j"
+#  Neo4j Sweden AB [https://neo4j.com]
+#  #
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  #
+#      https://www.apache.org/licenses/LICENSE-2.0
+#  #
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  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 unittest.mock import AsyncMock
+
+import pytest
+
+from neo4j_graphrag.experimental.pipeline import Component
+from neo4j_graphrag.experimental.pipeline.types.context import RunContext
+from .components import ComponentMultiply, ComponentMultiplyWithContext, IntResultModel
+
+
+def test_component_inputs() -> None:
+    inputs = ComponentMultiply.component_inputs
+    assert "number1" in inputs
+    assert inputs["number1"]["has_default"] is False
+    assert "number2" in inputs
+    assert inputs["number2"]["has_default"] is True
+
+
+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,
+        # we know this is a type and not a bool or str:
+        ComponentMultiply.component_outputs["result"]["annotation"],  # type: ignore
+    )
+
+
+@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,
+    )
+    # the type checker doesn't know about the type
+    # because the method is not re-declared
+    assert result.result == 2  # type: ignore
+
+
+@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()
+
+
+def test_component_missing_method() -> None:
+    with pytest.raises(RuntimeError) as e:
+
+        class WrongComponent(Component):
+            # we must have either run or run_with_context
+            pass
+
+    assert (
+        "You must implement either `run` or `run_with_context` in Component 'WrongComponent'"
+        in str(e)
+    )