Move tests from test_connect.py in test_pipeline.py and test_utils.py (…

…#7742)
deepset-ai · May 24, 2024 · 22289f5 · 22289f5
1 parent f5becf2
commit 22289f5
Show file tree

Hide file tree

Showing 3 changed files with 188 additions and 204 deletions.
diff --git a/test/core/pipeline/test_connect.py b/test/core/pipeline/test_connect.py
diff --git a/test/core/pipeline/test_pipeline.py b/test/core/pipeline/test_pipeline.py
@@ -15,7 +15,13 @@
 from haystack.components.routers import ConditionalRouter
 from haystack.core.component import component
 from haystack.core.component.types import InputSocket, OutputSocket, Variadic
-from haystack.core.errors import PipelineDrawingError, PipelineError, PipelineMaxLoops, PipelineRuntimeError
+from haystack.core.errors import (
+    PipelineConnectError,
+    PipelineDrawingError,
+    PipelineError,
+    PipelineMaxLoops,
+    PipelineRuntimeError,
+)
 from haystack.core.pipeline import Pipeline, PredefinedPipeline
 from haystack.core.serialization import DeserializationCallbacks
 from haystack.document_stores.in_memory import InMemoryDocumentStore
@@ -759,6 +765,178 @@ def test__prepare_component_input_data_with_non_existing_input(self, caplog):
             "please check your run parameters." in caplog.text
         )
 
+    def test_connect(self):
+        comp1 = component_class("Comp1", output_types={"value": int})()
+        comp2 = component_class("Comp2", input_types={"value": int})()
+        pipe = Pipeline()
+        pipe.add_component("comp1", comp1)
+        pipe.add_component("comp2", comp2)
+        assert pipe.connect("comp1.value", "comp2.value") is pipe
+
+        assert comp1.__haystack_output__.value.receivers == ["comp2"]
+        assert comp2.__haystack_input__.value.senders == ["comp1"]
+        assert list(pipe.graph.edges) == [("comp1", "comp2", "value/value")]
+
+    def test_connect_already_connected(self):
+        comp1 = component_class("Comp1", output_types={"value": int})()
+        comp2 = component_class("Comp2", input_types={"value": int})()
+        pipe = Pipeline()
+        pipe.add_component("comp1", comp1)
+        pipe.add_component("comp2", comp2)
+        pipe.connect("comp1.value", "comp2.value")
+        pipe.connect("comp1.value", "comp2.value")
+
+        assert comp1.__haystack_output__.value.receivers == ["comp2"]
+        assert comp2.__haystack_input__.value.senders == ["comp1"]
+        assert list(pipe.graph.edges) == [("comp1", "comp2", "value/value")]
+
+    def test_connect_with_sender_component_name(self):
+        comp1 = component_class("Comp1", output_types={"value": int})()
+        comp2 = component_class("Comp2", input_types={"value": int})()
+        pipe = Pipeline()
+        pipe.add_component("comp1", comp1)
+        pipe.add_component("comp2", comp2)
+        pipe.connect("comp1", "comp2.value")
+
+        assert comp1.__haystack_output__.value.receivers == ["comp2"]
+        assert comp2.__haystack_input__.value.senders == ["comp1"]
+        assert list(pipe.graph.edges) == [("comp1", "comp2", "value/value")]
+
+    def test_connect_with_receiver_component_name(self):
+        comp1 = component_class("Comp1", output_types={"value": int})()
+        comp2 = component_class("Comp2", input_types={"value": int})()
+        pipe = Pipeline()
+        pipe.add_component("comp1", comp1)
+        pipe.add_component("comp2", comp2)
+        pipe.connect("comp1.value", "comp2")
+
+        assert comp1.__haystack_output__.value.receivers == ["comp2"]
+        assert comp2.__haystack_input__.value.senders == ["comp1"]
+        assert list(pipe.graph.edges) == [("comp1", "comp2", "value/value")]
+
+    def test_connect_with_sender_and_receiver_component_name(self):
+        comp1 = component_class("Comp1", output_types={"value": int})()
+        comp2 = component_class("Comp2", input_types={"value": int})()
+        pipe = Pipeline()
+        pipe.add_component("comp1", comp1)
+        pipe.add_component("comp2", comp2)
+        pipe.connect("comp1", "comp2")
+
+        assert comp1.__haystack_output__.value.receivers == ["comp2"]
+        assert comp2.__haystack_input__.value.senders == ["comp1"]
+        assert list(pipe.graph.edges) == [("comp1", "comp2", "value/value")]
+
+    def test_connect_with_sender_not_in_pipeline(self):
+        comp2 = component_class("Comp2", input_types={"value": int})()
+        pipe = Pipeline()
+        pipe.add_component("comp2", comp2)
+        with pytest.raises(ValueError):
+            pipe.connect("comp1.value", "comp2.value")
+
+    def test_connect_with_receiver_not_in_pipeline(self):
+        comp1 = component_class("Comp1", output_types={"value": int})()
+        pipe = Pipeline()
+        pipe.add_component("comp1", comp1)
+        with pytest.raises(ValueError):
+            pipe.connect("comp1.value", "comp2.value")
+
+    def test_connect_with_sender_socket_name_not_in_pipeline(self):
+        comp1 = component_class("Comp1", output_types={"value": int})()
+        comp2 = component_class("Comp2", input_types={"value": int})()
+        pipe = Pipeline()
+        pipe.add_component("comp1", comp1)
+        pipe.add_component("comp2", comp2)
+        with pytest.raises(PipelineConnectError):
+            pipe.connect("comp1.non_existing", "comp2.value")
+
+    def test_connect_with_receiver_socket_name_not_in_pipeline(self):
+        comp1 = component_class("Comp1", output_types={"value": int})()
+        comp2 = component_class("Comp2", input_types={"value": int})()
+        pipe = Pipeline()
+        pipe.add_component("comp1", comp1)
+        pipe.add_component("comp2", comp2)
+        with pytest.raises(PipelineConnectError):
+            pipe.connect("comp1.value", "comp2.non_existing")
+
+    def test_connect_with_no_matching_types_and_same_names(self):
+        comp1 = component_class("Comp1", output_types={"value": int})()
+        comp2 = component_class("Comp2", input_types={"value": str})()
+        pipe = Pipeline()
+        pipe.add_component("comp1", comp1)
+        pipe.add_component("comp2", comp2)
+        with pytest.raises(PipelineConnectError):
+            pipe.connect("comp1", "comp2")
+
+    def test_connect_with_multiple_sender_connections_with_same_type_and_differing_name(self):
+        comp1 = component_class("Comp1", output_types={"val1": int, "val2": int})()
+        comp2 = component_class("Comp2", input_types={"value": int})()
+        pipe = Pipeline()
+        pipe.add_component("comp1", comp1)
+        pipe.add_component("comp2", comp2)
+        with pytest.raises(PipelineConnectError):
+            pipe.connect("comp1", "comp2")
+
+    def test_connect_with_multiple_receiver_connections_with_same_type_and_differing_name(self):
+        comp1 = component_class("Comp1", output_types={"value": int})()
+        comp2 = component_class("Comp2", input_types={"val1": int, "val2": int})()
+        pipe = Pipeline()
+        pipe.add_component("comp1", comp1)
+        pipe.add_component("comp2", comp2)
+        with pytest.raises(PipelineConnectError):
+            pipe.connect("comp1", "comp2")
+
+    def test_connect_with_multiple_sender_connections_with_same_type_and_same_name(self):
+        comp1 = component_class("Comp1", output_types={"value": int, "other": int})()
+        comp2 = component_class("Comp2", input_types={"value": int})()
+        pipe = Pipeline()
+        pipe.add_component("comp1", comp1)
+        pipe.add_component("comp2", comp2)
+        pipe.connect("comp1", "comp2")
+
+        assert comp1.__haystack_output__.value.receivers == ["comp2"]
+        assert comp2.__haystack_input__.value.senders == ["comp1"]
+        assert list(pipe.graph.edges) == [("comp1", "comp2", "value/value")]
+
+    def test_connect_with_multiple_receiver_connections_with_same_type_and_same_name(self):
+        comp1 = component_class("Comp1", output_types={"value": int})()
+        comp2 = component_class("Comp2", input_types={"value": int, "other": int})()
+        pipe = Pipeline()
+        pipe.add_component("comp1", comp1)
+        pipe.add_component("comp2", comp2)
+        pipe.connect("comp1", "comp2")
+
+        assert comp1.__haystack_output__.value.receivers == ["comp2"]
+        assert comp2.__haystack_input__.value.senders == ["comp1"]
+        assert list(pipe.graph.edges) == [("comp1", "comp2", "value/value")]
+
+    def test_connect_multiple_outputs_to_non_variadic_input(self):
+        comp1 = component_class("Comp1", output_types={"value": int})()
+        comp2 = component_class("Comp2", output_types={"value": int})()
+        comp3 = component_class("Comp3", input_types={"value": int})()
+        pipe = Pipeline()
+        pipe.add_component("comp1", comp1)
+        pipe.add_component("comp2", comp2)
+        pipe.add_component("comp3", comp3)
+        pipe.connect("comp1.value", "comp3.value")
+        with pytest.raises(PipelineConnectError):
+            pipe.connect("comp2.value", "comp3.value")
+
+    def test_connect_multiple_outputs_to_variadic_input(self):
+        comp1 = component_class("Comp1", output_types={"value": int})()
+        comp2 = component_class("Comp2", output_types={"value": int})()
+        comp3 = component_class("Comp3", input_types={"value": Variadic[int]})()
+        pipe = Pipeline()
+        pipe.add_component("comp1", comp1)
+        pipe.add_component("comp2", comp2)
+        pipe.add_component("comp3", comp3)
+        pipe.connect("comp1.value", "comp3.value")
+        pipe.connect("comp2.value", "comp3.value")
+
+        assert comp1.__haystack_output__.value.receivers == ["comp3"]
+        assert comp2.__haystack_output__.value.receivers == ["comp3"]
+        assert comp3.__haystack_input__.value.senders == ["comp1", "comp2"]
+        assert list(pipe.graph.edges) == [("comp1", "comp3", "value/value"), ("comp2", "comp3", "value/value")]
+
 
 @component
 class FakeComponent: