From 82bd12413965a2d2c5f17b9467ec08d6ed6efa5d Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Thu, 23 Nov 2023 08:39:45 -0300 Subject: [PATCH] Fix import order in test_custom_component.py --- tests/test_custom_component.py | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/tests/test_custom_component.py b/tests/test_custom_component.py index c60247668..b07753b8d 100644 --- a/tests/test_custom_component.py +++ b/tests/test_custom_component.py @@ -1,18 +1,17 @@ import ast -import pytest import types from uuid import uuid4 - +import pytest 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.code_parser import CodeParser, CodeSyntaxError from langflow.interface.custom.component import ( Component, ComponentCodeNullError, ) -from langflow.interface.custom.code_parser import CodeParser, CodeSyntaxError - +from langflow.services.database.models.flow import Flow, FlowCreate code_default = """ from langflow import Prompt @@ -229,9 +228,11 @@ def test_custom_component_get_function_entrypoint_return_type(): Test the get_function_entrypoint_return_type property of the CustomComponent class. """ + from langchain.schema import Document + custom_component = CustomComponent(code=code_default, function_entrypoint_name="build") 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(): @@ -414,7 +415,7 @@ class MyClass(CustomComponent): custom_component = CustomComponent(code=my_code, function_entrypoint_name="build") 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():