langflow/src/backend/tests/unit/test_custom_component.py

import ast
import types
from textwrap import dedent
from uuid import uuid4

import pytest
from langchain_core.documents import Document

from langflow.custom import Component, CustomComponent
from langflow.custom.code_parser.code_parser import CodeParser, CodeSyntaxError
from langflow.custom.custom_component.base_component import BaseComponent, ComponentCodeNullError
from langflow.custom.utils import build_custom_component_template
from langflow.services.database.models.flow import FlowCreate
from langflow.services.settings.feature_flags import FEATURE_FLAGS


@pytest.fixture
def client():
    pass


@pytest.fixture
def code_component_with_multiple_outputs():
    with open("src/backend/tests/data/component_multiple_outputs.py") as f:
        code = f.read()
        return Component(_code=code)


code_default = """
from langflow.custom import CustomComponent

from langflow.field_typing import BaseLanguageModel
from langchain.chains import LLMChain
from langchain.prompts import PromptTemplate
from langchain_core.documents import Document

import requests

class YourComponent(CustomComponent):
    display_name: str = "Your Component"
    description: str = "Your description"
    field_config = { "url": { "multiline": True, "required": True } }

    def build(self, url: str, llm: BaseLanguageModel) -> Document:
        return Document(page_content="Hello World")
"""


def test_code_parser_init():
    """
    Test the initialization of the CodeParser class.
    """
    parser = CodeParser(code_default)
    assert parser.code == code_default


def test_code_parser_get_tree():
    """
    Test the __get_tree method of the CodeParser class.
    """
    parser = CodeParser(code_default)
    tree = parser.get_tree()
    assert isinstance(tree, ast.AST)


def test_code_parser_syntax_error():
    """
    Test the __get_tree method raises the
    CodeSyntaxError when given incorrect syntax.
    """
    code_syntax_error = "zzz import os"

    parser = CodeParser(code_syntax_error)
    with pytest.raises(CodeSyntaxError):
        parser.get_tree()


def test_component_init():
    """
    Test the initialization of the Component class.
    """
    component = BaseComponent(_code=code_default, _function_entrypoint_name="build")
    assert component._code == code_default
    assert component._function_entrypoint_name == "build"


def test_component_get_code_tree():
    """
    Test the get_code_tree method of the Component class.
    """
    component = BaseComponent(_code=code_default, _function_entrypoint_name="build")
    tree = component.get_code_tree(component._code)
    assert "imports" in tree


def test_component_code_null_error():
    """
    Test the get_function method raises the
    ComponentCodeNullError when the code is empty.
    """
    component = BaseComponent(_code="", _function_entrypoint_name="")
    with pytest.raises(ComponentCodeNullError):
        component.get_function()


def test_custom_component_init():
    """
    Test the initialization of the CustomComponent class.
    """
    function_entrypoint_name = "build"

    custom_component = CustomComponent(_code=code_default, _function_entrypoint_name=function_entrypoint_name)
    assert custom_component._code == code_default
    assert custom_component._function_entrypoint_name == function_entrypoint_name


def test_custom_component_build_template_config():
    """
    Test the build_template_config property of the CustomComponent class.
    """
    custom_component = CustomComponent(_code=code_default, _function_entrypoint_name="build")
    config = custom_component.build_template_config()
    assert isinstance(config, dict)


def test_custom_component_get_function():
    """
    Test the get_function property of the CustomComponent class.
    """
    custom_component = CustomComponent(_code="def build(): pass", _function_entrypoint_name="build")
    my_function = custom_component.get_function()
    assert isinstance(my_function, types.FunctionType)


def test_code_parser_parse_imports_import():
    """
    Test the parse_imports method of the CodeParser
    class with an import statement.
    """
    parser = CodeParser(code_default)
    tree = parser.get_tree()
    for node in ast.walk(tree):
        if isinstance(node, ast.Import):
            parser.parse_imports(node)
    assert "requests" in parser.data["imports"]


def test_code_parser_parse_imports_importfrom():
    """
    Test the parse_imports method of the CodeParser
    class with an import from statement.
    """
    parser = CodeParser("from os import path")
    tree = parser.get_tree()
    for node in ast.walk(tree):
        if isinstance(node, ast.ImportFrom):
            parser.parse_imports(node)
    assert ("os", "path") in parser.data["imports"]


def test_code_parser_parse_functions():
    """
    Test the parse_functions method of the CodeParser class.
    """
    parser = CodeParser("def test(): pass")
    tree = parser.get_tree()
    for node in ast.walk(tree):
        if isinstance(node, ast.FunctionDef):
            parser.parse_functions(node)
    assert len(parser.data["functions"]) == 1
    assert parser.data["functions"][0]["name"] == "test"


def test_code_parser_parse_classes():
    """
    Test the parse_classes method of the CodeParser class.
    """
    parser = CodeParser("from langflow.custom import Component\n\nclass Test(Component): pass")
    tree = parser.get_tree()
    for node in ast.walk(tree):
        if isinstance(node, ast.ClassDef):
            parser.parse_classes(node)
    assert len(parser.data["classes"]) == 1
    assert parser.data["classes"][0]["name"] == "Test"


def test_code_parser_parse_classes_raises():
    """
    Test the parse_classes method of the CodeParser class.
    """
    parser = CodeParser("class Test: pass")
    tree = parser.get_tree()
    with pytest.raises(TypeError):
        for node in ast.walk(tree):
            if isinstance(node, ast.ClassDef):
                parser.parse_classes(node)


def test_code_parser_parse_global_vars():
    """
    Test the parse_global_vars method of the CodeParser class.
    """
    parser = CodeParser("x = 1")
    tree = parser.get_tree()
    for node in ast.walk(tree):
        if isinstance(node, ast.Assign):
            parser.parse_global_vars(node)
    assert len(parser.data["global_vars"]) == 1
    assert parser.data["global_vars"][0]["targets"] == ["x"]


def test_component_get_function_valid():
    """
    Test the get_function method of the Component
    class with valid code and function_entrypoint_name.
    """
    component = BaseComponent(_code="def build(): pass", _function_entrypoint_name="build")
    my_function = component.get_function()
    assert callable(my_function)


def test_custom_component_get_function_entrypoint_args():
    """
    Test the get_function_entrypoint_args
    property of the CustomComponent class.
    """
    custom_component = CustomComponent(_code=code_default, _function_entrypoint_name="build")
    args = custom_component.get_function_entrypoint_args
    assert len(args) == 3
    assert args[0]["name"] == "self"
    assert args[1]["name"] == "url"
    assert args[2]["name"] == "llm"


def test_custom_component_get_function_entrypoint_return_type():
    """
    Test the get_function_entrypoint_return_type
    property of the CustomComponent class.
    """

    custom_component = CustomComponent(_code=code_default, _function_entrypoint_name="build")
    return_type = custom_component.get_function_entrypoint_return_type
    assert return_type == [Document]


def test_custom_component_get_main_class_name():
    """
    Test the get_main_class_name property of the CustomComponent class.
    """
    custom_component = CustomComponent(_code=code_default, _function_entrypoint_name="build")
    class_name = custom_component.get_main_class_name
    assert class_name == "YourComponent"


def test_custom_component_get_function_valid():
    """
    Test the get_function property of the CustomComponent
    class with valid code and function_entrypoint_name.
    """
    custom_component = CustomComponent(_code="def build(): pass", _function_entrypoint_name="build")
    my_function = custom_component.get_function
    assert callable(my_function)


def test_code_parser_parse_arg_no_annotation():
    """
    Test the parse_arg method of the CodeParser class without an annotation.
    """
    parser = CodeParser("")
    arg = ast.arg(arg="x", annotation=None)
    result = parser.parse_arg(arg, None)
    assert result["name"] == "x"
    assert "type" not in result


def test_code_parser_parse_arg_with_annotation():
    """
    Test the parse_arg method of the CodeParser class with an annotation.
    """
    parser = CodeParser("")
    arg = ast.arg(arg="x", annotation=ast.Name(id="int", ctx=ast.Load()))
    result = parser.parse_arg(arg, None)
    assert result["name"] == "x"
    assert result["type"] == "int"


def test_code_parser_parse_callable_details_no_args():
    """
    Test the parse_callable_details method of the
    CodeParser class with a function with no arguments.
    """
    parser = CodeParser("")
    node = ast.FunctionDef(
        name="test",
        args=ast.arguments(args=[], vararg=None, kwonlyargs=[], kw_defaults=[], kwarg=None, defaults=[]),
        body=[],
        decorator_list=[],
        returns=None,
    )
    result = parser.parse_callable_details(node)
    assert result["name"] == "test"
    assert len(result["args"]) == 0


def test_code_parser_parse_assign():
    """
    Test the parse_assign method of the CodeParser class.
    """
    parser = CodeParser("")
    stmt = ast.Assign(targets=[ast.Name(id="x", ctx=ast.Store())], value=ast.Num(n=1))
    result = parser.parse_assign(stmt)
    assert result["name"] == "x"
    assert result["value"] == "1"


def test_code_parser_parse_ann_assign():
    """
    Test the parse_ann_assign method of the CodeParser class.
    """
    parser = CodeParser("")
    stmt = ast.AnnAssign(
        target=ast.Name(id="x", ctx=ast.Store()),
        annotation=ast.Name(id="int", ctx=ast.Load()),
        value=ast.Num(n=1),
        simple=1,
    )
    result = parser.parse_ann_assign(stmt)
    assert result["name"] == "x"
    assert result["value"] == "1"
    assert result["annotation"] == "int"


def test_code_parser_parse_function_def_not_init():
    """
    Test the parse_function_def method of the
    CodeParser class with a function that is not __init__.
    """
    parser = CodeParser("")
    stmt = ast.FunctionDef(
        name="test",
        args=ast.arguments(args=[], vararg=None, kwonlyargs=[], kw_defaults=[], kwarg=None, defaults=[]),
        body=[],
        decorator_list=[],
        returns=None,
    )
    result, is_init = parser.parse_function_def(stmt)
    assert result["name"] == "test"
    assert not is_init


def test_code_parser_parse_function_def_init():
    """
    Test the parse_function_def method of the
    CodeParser class with an __init__ function.
    """
    parser = CodeParser("")
    stmt = ast.FunctionDef(
        name="__init__",
        args=ast.arguments(args=[], vararg=None, kwonlyargs=[], kw_defaults=[], kwarg=None, defaults=[]),
        body=[],
        decorator_list=[],
        returns=None,
    )
    result, is_init = parser.parse_function_def(stmt)
    assert result["name"] == "__init__"
    assert is_init


def test_component_get_code_tree_syntax_error():
    """
    Test the get_code_tree method of the Component class
    raises the CodeSyntaxError when given incorrect syntax.
    """
    component = BaseComponent(_code="import os as", _function_entrypoint_name="build")
    with pytest.raises(CodeSyntaxError):
        component.get_code_tree(component._code)


def test_custom_component_class_template_validation_no_code():
    """
    Test the _class_template_validation method of the CustomComponent class
    raises the HTTPException when the code is None.
    """
    custom_component = CustomComponent(_code=None, _function_entrypoint_name="build")
    with pytest.raises(TypeError):
        custom_component.get_function()


def test_custom_component_get_code_tree_syntax_error():
    """
    Test the get_code_tree method of the CustomComponent class
    raises the CodeSyntaxError when given incorrect syntax.
    """
    custom_component = CustomComponent(_code="import os as", _function_entrypoint_name="build")
    with pytest.raises(CodeSyntaxError):
        custom_component.get_code_tree(custom_component._code)


def test_custom_component_get_function_entrypoint_args_no_args():
    """
    Test the get_function_entrypoint_args property of
    the CustomComponent class with a build method with no arguments.
    """
    my_code = """
from langflow.custom import CustomComponent
class MyMainClass(CustomComponent):
    def build():
        pass"""

    custom_component = CustomComponent(_code=my_code, _function_entrypoint_name="build")
    args = custom_component.get_function_entrypoint_args
    assert len(args) == 0


def test_custom_component_get_function_entrypoint_return_type_no_return_type():
    """
    Test the get_function_entrypoint_return_type property of the
    CustomComponent class with a build method with no return type.
    """
    my_code = """
from langflow.custom import CustomComponent
class MyClass(CustomComponent):
    def build():
        pass"""

    custom_component = CustomComponent(_code=my_code, _function_entrypoint_name="build")
    return_type = custom_component.get_function_entrypoint_return_type
    assert return_type == []


def test_custom_component_get_main_class_name_no_main_class():
    """
    Test the get_main_class_name property of the
    CustomComponent class when there is no main class.
    """
    my_code = """
def build():
    pass"""

    custom_component = CustomComponent(_code=my_code, _function_entrypoint_name="build")
    class_name = custom_component.get_main_class_name
    assert class_name == ""


def test_custom_component_build_not_implemented():
    """
    Test the build method of the CustomComponent
    class raises the NotImplementedError.
    """
    custom_component = CustomComponent(_code="def build(): pass", _function_entrypoint_name="build")
    with pytest.raises(NotImplementedError):
        custom_component.build()


def test_build_config_no_code():
    component = CustomComponent(_code=None)

    assert component.get_function_entrypoint_args == []
    assert component.get_function_entrypoint_return_type == []


@pytest.fixture
def component(client, active_user):
    return CustomComponent(
        user_id=active_user.id,
        field_config={
            "fields": {
                "llm": {"type": "str"},
                "url": {"type": "str"},
                "year": {"type": "int"},
            }
        },
    )


@pytest.fixture(scope="session")
def test_flow(db):
    flow_data = {
        "nodes": [{"id": "1"}, {"id": "2"}],
        "edges": [{"source": "1", "target": "2"}],
    }

    # Create flow
    flow = FlowCreate(id=uuid4(), name="Test Flow", description="Fixture flow", data=flow_data)

    # Add to database
    db.add(flow)
    db.commit()

    yield flow

    # Clean up
    db.delete(flow)
    db.commit()


@pytest.fixture(scope="session")
def db(app):
    # Setup database for tests
    yield app.db

    # Teardown
    app.db.drop_all()


def test_list_flows_return_type(component):
    flows = component.list_flows()
    assert isinstance(flows, list)


def test_build_config_return_type(component):
    config = component.build_config()
    assert isinstance(config, dict)


def test_build_config_has_fields(component):
    config = component.build_config()
    assert "fields" in config


def test_build_config_fields_dict(component):
    config = component.build_config()
    assert isinstance(config["fields"], dict)


def test_build_config_field_keys(component):
    config = component.build_config()
    assert all(isinstance(key, str) for key in config["fields"])


def test_build_config_field_values_dict(component):
    config = component.build_config()
    assert all(isinstance(value, dict) for value in config["fields"].values())


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"]


def test_feature_flags_add_toolkit_output(active_user, code_component_with_multiple_outputs):
    frontnd_node_dict, _ = build_custom_component_template(code_component_with_multiple_outputs, active_user.id)
    len_outputs = len(frontnd_node_dict["outputs"])
    FEATURE_FLAGS.add_toolkit_output = True
    frontnd_node_dict, _ = build_custom_component_template(code_component_with_multiple_outputs, active_user.id)
    assert len(frontnd_node_dict["outputs"]) == len_outputs + 1


def test_custom_component_subclass_from_lctoolcomponent():
    # Import LCToolComponent and create a subclass
    code = dedent("""
    from langflow.base.langchain_utilities.model import LCToolComponent
    from langchain_core.tools import Tool
    class MyComponent(LCToolComponent):
        name: str = "MyComponent"
        description: str = "MyComponent"

        def build_tool(self) -> Tool:
            return Tool(name="MyTool", description="MyTool")

        def run_model(self)-> Data:
            return Data(data="Hello World")
    """)
    component = Component(_code=code)
    frontend_node, _ = build_custom_component_template(component)
    assert "outputs" in frontend_node
    assert frontend_node["outputs"][0]["types"] != []
    assert frontend_node["outputs"][1]["types"] != []