Skip to content

Commit

Permalink
🐛 fix(custom_component.py): handle case when return_type is None to p…
Browse files Browse the repository at this point in the history
…revent errors

🐛 fix(test_custom_component.py): update assertion to expect return_type as a list instead of a string
🐛 fix(test_vectorstore_template.py): update assertion to check if all vectorstores in settings are present in the response
  • Loading branch information
ogabrielluiz committed Aug 9, 2023
1 parent 30aad35 commit 8c41415
Show file tree
Hide file tree
Showing 3 changed files with 7 additions and 4 deletions.
2 changes: 2 additions & 0 deletions src/backend/langflow/interface/custom/custom_component.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 []
Expand Down
6 changes: 3 additions & 3 deletions tests/test_custom_component.py
Original file line number Diff line number Diff line change
Expand Up @@ -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():
Expand Down Expand Up @@ -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():
Expand Down Expand Up @@ -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
Expand Down
3 changes: 2 additions & 1 deletion tests/test_vectorstore_template.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)

0 comments on commit 8c41415

Please sign in to comment.