From 8c41415a322d2b4bc3ac7ba77767da3a5660b055 Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Wed, 9 Aug 2023 17:50:03 -0300 Subject: [PATCH] =?UTF-8?q?=F0=9F=90=9B=20fix(custom=5Fcomponent.py):=20ha?= =?UTF-8?q?ndle=20case=20when=20return=5Ftype=20is=20None=20to=20prevent?= =?UTF-8?q?=20errors=20=F0=9F=90=9B=20fix(test=5Fcustom=5Fcomponent.py):?= =?UTF-8?q?=20update=20assertion=20to=20expect=20return=5Ftype=20as=20a=20?= =?UTF-8?q?list=20instead=20of=20a=20string=20=F0=9F=90=9B=20fix(test=5Fve?= =?UTF-8?q?ctorstore=5Ftemplate.py):=20update=20assertion=20to=20check=20i?= =?UTF-8?q?f=20all=20vectorstores=20in=20settings=20are=20present=20in=20t?= =?UTF-8?q?he=20response?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/backend/langflow/interface/custom/custom_component.py | 2 ++ tests/test_custom_component.py | 6 +++--- tests/test_vectorstore_template.py | 3 ++- 3 files changed, 7 insertions(+), 4 deletions(-) diff --git a/src/backend/langflow/interface/custom/custom_component.py b/src/backend/langflow/interface/custom/custom_component.py index d7dd280e647..c1de48a98fd 100644 --- a/src/backend/langflow/interface/custom/custom_component.py +++ b/src/backend/langflow/interface/custom/custom_component.py @@ -118,6 +118,8 @@ def get_function_entrypoint_return_type(self) -> List[str]: build_method = build_methods[0] return_type = build_method["return_type"] + if not return_type: + return [] # If the return type is not a Union, then we just return it as a list if "Union" not in return_type: return [return_type] if return_type in self.return_type_valid_list else [] diff --git a/tests/test_custom_component.py b/tests/test_custom_component.py index 199906dda88..01fe8b597ef 100644 --- a/tests/test_custom_component.py +++ b/tests/test_custom_component.py @@ -241,7 +241,7 @@ def test_custom_component_get_function_entrypoint_return_type(): 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(): @@ -436,7 +436,7 @@ def build(): custom_component = CustomComponent(code=my_code, function_entrypoint_name="build") return_type = custom_component.get_function_entrypoint_return_type - assert return_type is None + assert return_type == [] def test_custom_component_get_main_class_name_no_main_class(): @@ -469,7 +469,7 @@ def test_build_config_no_code(): component = CustomComponent(code=None) assert component.get_function_entrypoint_args == "" - assert component.get_function_entrypoint_return_type == "" + assert component.get_function_entrypoint_return_type == [] @pytest.fixture diff --git a/tests/test_vectorstore_template.py b/tests/test_vectorstore_template.py index bac950ee116..1161d498271 100644 --- a/tests/test_vectorstore_template.py +++ b/tests/test_vectorstore_template.py @@ -9,4 +9,5 @@ def test_vectorstores_settings(client: TestClient): assert response.status_code == 200 json_response = response.json() vectorstores = json_response["vectorstores"] - assert set(vectorstores.keys()) == set(settings.VECTORSTORES) + settings_vecs = set(settings.VECTORSTORES) + assert all(vs in vectorstores for vs in settings_vecs)