diff --git a/tests/data/component_multiple_outputs.py b/tests/data/component_multiple_outputs.py new file mode 100644 index 000000000..7a01fafba --- /dev/null +++ b/tests/data/component_multiple_outputs.py @@ -0,0 +1,19 @@ +from langflow.custom import CustomComponent +from langflow.template.field.base import Input, Output + + +class MultipleOutputsComponent(CustomComponent): + inputs = [ + Input(display_name="Input", name="input", field_type=str), + Input(display_name="Number", name="number", field_type=int), + ] + outputs = [ + Output(display_name="Certain Output", method="certain_output", name="certain_output"), + Output(name="Other Output", method="other_output"), + ] + + def certain_output(self) -> str: + return f"This is my string input: {self.input}" + + def other_output(self) -> int: + return f"This is my int input multiplied by 2: {self.number * 2}" diff --git a/tests/test_custom_component.py b/tests/test_custom_component.py index 0b64f2c1a..461d3c6eb 100644 --- a/tests/test_custom_component.py +++ b/tests/test_custom_component.py @@ -8,8 +8,17 @@ from langchain_core.documents import Document from langflow.custom import CustomComponent from langflow.custom.code_parser.code_parser import CodeParser, CodeSyntaxError from langflow.custom.custom_component.component import Component, ComponentCodeNullError +from langflow.custom.utils import build_custom_component_template from langflow.services.database.models.flow import Flow, FlowCreate + +@pytest.fixture +def code_component_with_multiple_outputs(): + with open("tests/data/component_multiple_outputs.py", "r") as f: + code = f.read() + return CustomComponent(code=code) + + code_default = """ from langflow.field_typing import Prompt from langflow.custom import CustomComponent @@ -517,3 +526,8 @@ def test_build_config_field_value_keys(component): config = component.build_config() field_values = config["fields"].values() assert all("type" in value for value in field_values) + + +def test_custom_component_multiple_outputs(code_component_with_multiple_outputs, active_user): + frontnd_node_dict, _ = build_custom_component_template(code_component_with_multiple_outputs, active_user.id) + assert frontnd_node_dict["outputs"][0]["types"] == ["Text"]