Fix import order in test_custom_component.py

This commit is contained in:
Gabriel Luiz Freitas Almeida 2023-11-23 08:39:45 -03:00
commit 82bd124139

View file

@ -1,18 +1,17 @@
import ast import ast
import pytest
import types import types
from uuid import uuid4 from uuid import uuid4
import pytest
from fastapi import HTTPException from fastapi import HTTPException
from langflow.services.database.models.flow import Flow, FlowCreate from langflow.field_typing.constants import Data
from langflow.interface.custom.base import CustomComponent from langflow.interface.custom.base import CustomComponent
from langflow.interface.custom.code_parser import CodeParser, CodeSyntaxError
from langflow.interface.custom.component import ( from langflow.interface.custom.component import (
Component, Component,
ComponentCodeNullError, ComponentCodeNullError,
) )
from langflow.interface.custom.code_parser import CodeParser, CodeSyntaxError from langflow.services.database.models.flow import Flow, FlowCreate
code_default = """ code_default = """
from langflow import Prompt from langflow import Prompt
@ -229,9 +228,11 @@ def test_custom_component_get_function_entrypoint_return_type():
Test the get_function_entrypoint_return_type Test the get_function_entrypoint_return_type
property of the CustomComponent class. property of the CustomComponent class.
""" """
from langchain.schema import Document
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 return_type = custom_component.get_function_entrypoint_return_type
assert return_type == ["Document"] assert return_type == [Document]
def test_custom_component_get_main_class_name(): def test_custom_component_get_main_class_name():
@ -414,7 +415,7 @@ class MyClass(CustomComponent):
custom_component = CustomComponent(code=my_code, function_entrypoint_name="build") custom_component = CustomComponent(code=my_code, function_entrypoint_name="build")
return_type = custom_component.get_function_entrypoint_return_type return_type = custom_component.get_function_entrypoint_return_type
assert return_type == [] assert return_type == [Data]
def test_custom_component_get_main_class_name_no_main_class(): def test_custom_component_get_main_class_name_no_main_class():