src/backend/tests/unit/custom/custom_component/test_update_outputs.py

import pytest
from langflow.base.tools.constants import TOOL_OUTPUT_DISPLAY_NAME, TOOL_OUTPUT_NAME
from langflow.custom.custom_component.component import Component


class TestComponentOutputs:
    def test_run_and_validate_update_outputs_tool_mode(self):
        """Test run_and_validate_update_outputs with tool_mode field."""

        class TestComponent(Component):
            def build(self) -> None:
                pass

        component = TestComponent()

        # Create a frontend node with regular outputs
        original_outputs = [
            {
                "name": "regular_output",
                "type": "str",
                "display_name": "Regular Output",
                "method": "get_output",
                "types": ["Any"],
                "selected": "Any",
                "value": "__UNDEFINED__",
                "cache": True,
                "required_inputs": None,
                "hidden": None,
            }
        ]
        frontend_node = {
            "outputs": original_outputs.copy()  # Make a copy to preserve original
        }

        # Test enabling tool mode
        updated_node = component.run_and_validate_update_outputs(
            frontend_node=frontend_node.copy(),  # Use a copy to avoid modifying original
            field_name="tool_mode",
            field_value=True,
        )

        # Verify tool output is added and regular output is removed
        assert len(updated_node["outputs"]) == 1
        assert updated_node["outputs"][0]["name"] == TOOL_OUTPUT_NAME
        assert updated_node["outputs"][0]["display_name"] == TOOL_OUTPUT_DISPLAY_NAME

        # Test disabling tool mode - use the original frontend node
        updated_node = component.run_and_validate_update_outputs(
            frontend_node={"outputs": original_outputs.copy()},  # Use original outputs
            field_name="tool_mode",
            field_value=False,
        )

        # Verify original outputs are restored
        assert len(updated_node["outputs"]) == 1
        # Compare only essential fields instead of the entire dict
        assert updated_node["outputs"][0]["name"] == original_outputs[0]["name"]
        assert updated_node["outputs"][0]["display_name"] == original_outputs[0]["display_name"]
        assert updated_node["outputs"][0]["method"] == original_outputs[0]["method"]
        assert "types" in updated_node["outputs"][0]
        assert "selected" in updated_node["outputs"][0]

    def test_run_and_validate_update_outputs_invalid_output(self):
        """Test run_and_validate_update_outputs with invalid output structure."""

        class TestComponent(Component):
            def build(self) -> None:
                pass

        component = TestComponent()

        # Create a frontend node with invalid output structure
        frontend_node = {"outputs": [{"invalid_field": "value"}]}

        # Test validation fails for invalid output
        with pytest.raises(ValueError, match="Invalid output: 1 validation error for Output"):
            component.run_and_validate_update_outputs(
                frontend_node=frontend_node, field_name="some_field", field_value="some_value"
            )

    def test_run_and_validate_update_outputs_custom_update(self):
        """Test run_and_validate_update_outputs with custom update logic."""

        class CustomComponent(Component):
            def build(self) -> None:
                pass

            def get_custom(self) -> str:
                """Method that returns a string."""
                return "custom output"

            def update_outputs(self, frontend_node, field_name, field_value):  # noqa: ARG002
                if field_name == "custom_field":
                    frontend_node["outputs"].append(
                        {
                            "name": "custom_output",
                            "type": "str",
                            "display_name": "Custom Output",
                            "method": "get_custom",
                            "types": ["Any"],
                            "selected": "Any",
                            "value": "__UNDEFINED__",
                            "cache": True,
                            "required_inputs": None,
                            "hidden": None,
                        }
                    )
                return frontend_node

        component = CustomComponent()
        frontend_node = {"outputs": []}

        # Test custom update logic
        updated_node = component.run_and_validate_update_outputs(
            frontend_node=frontend_node, field_name="custom_field", field_value="custom_value"
        )

        assert len(updated_node["outputs"]) == 1
        assert updated_node["outputs"][0]["name"] == "custom_output"
        assert updated_node["outputs"][0]["display_name"] == "Custom Output"
        assert updated_node["outputs"][0]["method"] == "get_custom"
        assert "types" in updated_node["outputs"][0]
        assert "selected" in updated_node["outputs"][0]

    def test_run_and_validate_update_outputs_with_existing_tool_output(self):
        """Test run_and_validate_update_outputs when tool output already exists."""

        class TestComponent(Component):
            def build(self) -> None:
                pass

            def to_toolkit(self) -> list:
                """Method that returns a list of tools."""
                return []

        component = TestComponent()

        # Create a frontend node with tool output already present
        frontend_node = {
            "outputs": [
                {
                    "name": TOOL_OUTPUT_NAME,  # Use constant instead of hardcoded string
                    "type": "Tool",
                    "display_name": TOOL_OUTPUT_DISPLAY_NAME,  # Use constant
                    "method": "to_toolkit",
                    "types": ["Tool"],
                    "selected": "Tool",
                    "value": "__UNDEFINED__",
                    "cache": True,
                    "required_inputs": None,
                    "hidden": None,
                }
            ]
        }

        # Test enabling tool mode doesn't duplicate tool output
        updated_node = component.run_and_validate_update_outputs(
            frontend_node=frontend_node, field_name="tool_mode", field_value=True
        )

        assert len(updated_node["outputs"]) == 1
        assert updated_node["outputs"][0]["name"] == TOOL_OUTPUT_NAME  # Use constant
        assert updated_node["outputs"][0]["display_name"] == TOOL_OUTPUT_DISPLAY_NAME  # Use constant
        assert "types" in updated_node["outputs"][0]
        assert "selected" in updated_node["outputs"][0]

    def test_run_and_validate_update_outputs_with_multiple_outputs(self):
        """Test run_and_validate_update_outputs with multiple outputs."""

        class TestComponent(Component):
            def build(self) -> None:
                pass

            def get_output1(self) -> str:
                """Method that returns a string."""
                return "output1"

            def get_output2(self) -> str:
                """Method that returns a string."""
                return "output2"

            def update_outputs(self, frontend_node, field_name, field_value):  # noqa: ARG002
                if field_name == "add_output":
                    frontend_node["outputs"].extend(
                        [
                            {
                                "name": "output1",
                                "type": "str",
                                "display_name": "Output 1",
                                "method": "get_output1",
                            },
                            {
                                "name": "output2",
                                "type": "str",
                                "display_name": "Output 2",
                                "method": "get_output2",
                            },
                        ]
                    )
                return frontend_node

        component = TestComponent()
        frontend_node = {"outputs": []}

        # Test adding multiple outputs
        updated_node = component.run_and_validate_update_outputs(
            frontend_node=frontend_node, field_name="add_output", field_value=True
        )

        assert len(updated_node["outputs"]) == 2
        assert updated_node["outputs"][0]["name"] == "output1"
        assert updated_node["outputs"][1]["name"] == "output2"
        for output in updated_node["outputs"]:
            assert "types" in output
            assert "selected" in output
            # The component adds only 'Text' type for string outputs
            assert set(output["types"]) == {"Text"}
            assert output["selected"] == "Text"

    def test_run_and_validate_update_outputs_output_validation(self):
        """Test output validation in run_and_validate_update_outputs."""

        class TestComponent(Component):
            def build(self) -> None:
                pass

            def get_test(self) -> str:
                """Test method."""
                return "test"

        component = TestComponent()

        # Test invalid method name case
        invalid_node = {
            "outputs": [{"name": "test", "type": "str", "method": "nonexistent_method", "display_name": "Test"}]
        }

        with pytest.raises(AttributeError, match="nonexistent_method not found in TestComponent"):
            component.run_and_validate_update_outputs(frontend_node=invalid_node, field_name="test", field_value=True)

        # Test missing method case
        invalid_node = {"outputs": [{"name": "test", "type": "str", "display_name": "Test"}]}

        with pytest.raises(ValueError, match="Output test does not have a method"):
            component.run_and_validate_update_outputs(frontend_node=invalid_node, field_name="test", field_value=True)