Refactor tests and add new test file for data components

tests/test_custom_component.py

@@ -3,19 +3,21 @@ import types
 from uuid import uuid4
 
 import pytest
-from fastapi import HTTPException
 from langchain_core.documents import Document
+
 from langflow.interface.custom.base import CustomComponent
-from langflow.interface.custom.code_parser.code_parser import CodeParser, CodeSyntaxError
+from langflow.interface.custom.code_parser.code_parser import (
+    CodeParser,
+    CodeSyntaxError,
+)
 from langflow.interface.custom.custom_component.component import (
     Component,
     ComponentCodeNullError,
-    ComponentFunctionEntrypointNameNullError,
 )
 from langflow.services.database.models.flow import Flow, FlowCreate
 
 code_default = """
-from langflow import Prompt
+from langflow.field_typing import Prompt
 from langflow.interface.custom.custom_component import CustomComponent
 
 from langchain.llms.base import BaseLLM
@@ -96,23 +98,15 @@ def test_component_code_null_error():
         component.get_function()
 
 
-def test_component_function_entrypoint_name_null_error():
-    """
-    Test the get_function method raises the ComponentFunctionEntrypointNameNullError
-    when the function_entrypoint_name is empty.
-    """
-    component = Component(code=code_default, function_entrypoint_name="")
-    with pytest.raises(ComponentFunctionEntrypointNameNullError):
-        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)
+    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
 
@@ -121,7 +115,9 @@ 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")
+    custom_component = CustomComponent(
+        code=code_default, function_entrypoint_name="build"
+    )
     config = custom_component.build_template_config()
     assert isinstance(config, dict)
 
@@ -130,8 +126,10 @@ 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
+    custom_component = CustomComponent(
+        code="def build(): pass", function_entrypoint_name="build"
+    )
+    my_function = custom_component.get_function()
     assert isinstance(my_function, types.FunctionType)
 
 
@@ -215,7 +213,9 @@ 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")
+    custom_component = CustomComponent(
+        code=code_default, function_entrypoint_name="build"
+    )
     args = custom_component.get_function_entrypoint_args
     assert len(args) == 4
     assert args[0]["name"] == "self"
@@ -229,7 +229,9 @@ def test_custom_component_get_function_entrypoint_return_type():
     property of the CustomComponent class.
     """
 
-    custom_component = CustomComponent(code=code_default, function_entrypoint_name="build")
+    custom_component = CustomComponent(
+        code=code_default, function_entrypoint_name="build"
+    )
     return_type = custom_component.get_function_entrypoint_return_type
     assert return_type == [Document]
 
@@ -238,7 +240,9 @@ 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")
+    custom_component = CustomComponent(
+        code=code_default, function_entrypoint_name="build"
+    )
     class_name = custom_component.get_main_class_name
     assert class_name == "YourComponent"
 
@@ -248,7 +252,9 @@ 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")
+    custom_component = CustomComponent(
+        code="def build(): pass", function_entrypoint_name="build"
+    )
     my_function = custom_component.get_function
     assert callable(my_function)
 
@@ -283,7 +289,9 @@ def test_code_parser_parse_callable_details_no_args():
     parser = CodeParser("")
     node = ast.FunctionDef(
         name="test",
-        args=ast.arguments(args=[], vararg=None, kwonlyargs=[], kw_defaults=[], kwarg=None, defaults=[]),
+        args=ast.arguments(
+            args=[], vararg=None, kwonlyargs=[], kw_defaults=[], kwarg=None, defaults=[]
+        ),
         body=[],
         decorator_list=[],
         returns=None,
@@ -329,7 +337,9 @@ def test_code_parser_parse_function_def_not_init():
     parser = CodeParser("")
     stmt = ast.FunctionDef(
         name="test",
-        args=ast.arguments(args=[], vararg=None, kwonlyargs=[], kw_defaults=[], kwarg=None, defaults=[]),
+        args=ast.arguments(
+            args=[], vararg=None, kwonlyargs=[], kw_defaults=[], kwarg=None, defaults=[]
+        ),
         body=[],
         decorator_list=[],
         returns=None,
@@ -347,7 +357,9 @@ def test_code_parser_parse_function_def_init():
     parser = CodeParser("")
     stmt = ast.FunctionDef(
         name="__init__",
-        args=ast.arguments(args=[], vararg=None, kwonlyargs=[], kw_defaults=[], kwarg=None, defaults=[]),
+        args=ast.arguments(
+            args=[], vararg=None, kwonlyargs=[], kw_defaults=[], kwarg=None, defaults=[]
+        ),
         body=[],
         decorator_list=[],
         returns=None,
@@ -373,7 +385,7 @@ def test_custom_component_class_template_validation_no_code():
     raises the HTTPException when the code is None.
     """
     custom_component = CustomComponent(code=None, function_entrypoint_name="build")
-    with pytest.raises(HTTPException):
+    with pytest.raises(TypeError):
         custom_component.get_function()
 
 
@@ -382,7 +394,9 @@ 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")
+    custom_component = CustomComponent(
+        code="import os as", function_entrypoint_name="build"
+    )
     with pytest.raises(CodeSyntaxError):
         custom_component.get_code_tree(custom_component.code)
 
@@ -436,7 +450,9 @@ 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")
+    custom_component = CustomComponent(
+        code="def build(): pass", function_entrypoint_name="build"
+    )
     with pytest.raises(NotImplementedError):
         custom_component.build()
 
@@ -444,7 +460,7 @@ def test_custom_component_build_not_implemented():
 def test_build_config_no_code():
     component = CustomComponent(code=None)
 
-    assert component.get_function_entrypoint_args == ""
+    assert component.get_function_entrypoint_args == []
     assert component.get_function_entrypoint_return_type == []
 
 
@@ -470,7 +486,9 @@ def test_flow(db):
     }
 
     # Create flow
-    flow = FlowCreate(id=uuid4(), name="Test Flow", description="Fixture flow", data=flow_data)
+    flow = FlowCreate(
+        id=uuid4(), name="Test Flow", description="Fixture flow", data=flow_data
+    )
 
     # Add to database
     db.add(flow)