haystack/test/core/component/test_sockets.py
Silvano Cerza 0191b1e6e4
feat: Change Component's I/O dunder type (#6916)
* Add Pipeline.get_component_name() method

* Add utility class to ease discoverability of Component I/O

* Move InputOutput in component package

* Rename InputOutput to _InputOutput

* Raise if inputs or outputs field already exist

* Fix tests

* Add release notes

* Move InputSocket and OutputSocket in types package

* Move _InputOutput in socket package

* Rename _InputOutput class to Sockets

* Simplify Sockets class

* Dictch I/O dunder fields in favour of inputs and outputs fields

* Update Sockets docstrings

* Update release notes

* Fix mypy

* Remove unnecessary assignment

* Remove unused logging

* Change SocketsType to SocketsIOType to avoid confusion

* Change sockets type and name

* Change Sockets.__repr__ to return component instance

* Fix linting

* Fix sockets tests

* Revert to dunder fields for Component IO

* Use singular in IO dunder fields

* Delete release notes

* Update haystack/core/component/types.py

Co-authored-by: Massimiliano Pippi <mpippi@gmail.com>

---------

Co-authored-by: Massimiliano Pippi <mpippi@gmail.com>
2024-02-05 17:46:45 +01:00

58 lines
2.6 KiB
Python

import pytest
from haystack.core.component.sockets import InputSocket, Sockets
from haystack.core.pipeline import Pipeline
from haystack.testing.factory import component_class
class TestSockets:
def test_init(self):
comp = component_class("SomeComponent", input_types={"input_1": int, "input_2": int})()
sockets = {"input_1": InputSocket("input_1", int), "input_2": InputSocket("input_2", int)}
io = Sockets(component=comp, sockets_dict=sockets, sockets_io_type=InputSocket)
assert io._component == comp
assert "input_1" in io.__dict__
assert io.__dict__["input_1"] == comp.__haystack_input__._sockets_dict["input_1"]
assert "input_2" in io.__dict__
assert io.__dict__["input_2"] == comp.__haystack_input__._sockets_dict["input_2"]
def test_init_with_empty_sockets(self):
comp = component_class("SomeComponent")()
io = Sockets(component=comp, sockets_dict={}, sockets_io_type=InputSocket)
assert io._component == comp
assert io._sockets_dict == {}
def test_component_name(self):
comp = component_class("SomeComponent")()
io = Sockets(component=comp, sockets_dict={}, sockets_io_type=InputSocket)
assert io._component_name() == str(comp)
def test_component_name_added_to_pipeline(self):
comp = component_class("SomeComponent")()
pipeline = Pipeline()
pipeline.add_component("my_component", comp)
io = Sockets(component=comp, sockets_dict={}, sockets_io_type=InputSocket)
assert io._component_name() == "my_component"
def test_getattribute(self):
comp = component_class("SomeComponent", input_types={"input_1": int, "input_2": int})()
io = Sockets(component=comp, sockets_dict=comp.__haystack_input__._sockets_dict, sockets_io_type=InputSocket)
assert io.input_1 == comp.__haystack_input__._sockets_dict["input_1"]
assert io.input_2 == comp.__haystack_input__._sockets_dict["input_2"]
def test_getattribute_non_existing_socket(self):
comp = component_class("SomeComponent", input_types={"input_1": int, "input_2": int})()
io = Sockets(component=comp, sockets_dict=comp.__haystack_input__._sockets_dict, sockets_io_type=InputSocket)
with pytest.raises(AttributeError):
io.input_3
def test_repr(self):
comp = component_class("SomeComponent", input_types={"input_1": int, "input_2": int})()
io = Sockets(component=comp, sockets_dict=comp.__haystack_input__._sockets_dict, sockets_io_type=InputSocket)
res = repr(io)
assert res == f"{comp} inputs:\n - input_1: int\n - input_2: int"