Skip to content

Commit

Permalink
fix(node): correct file property name in function switch (langgenius#…
Browse files Browse the repository at this point in the history
  • Loading branch information
laipz8200 authored and BenjaminX committed Nov 5, 2024
1 parent 94add0d commit 07cd8d1
Show file tree
Hide file tree
Showing 2 changed files with 47 additions and 4 deletions.
2 changes: 1 addition & 1 deletion api/core/workflow/nodes/list_operator/node.py
Original file line number Diff line number Diff line change
Expand Up @@ -157,7 +157,7 @@ def _get_file_extract_string_func(*, key: str) -> Callable[[File], str]:
return lambda x: x.type
case "extension":
return lambda x: x.extension or ""
case "mimetype":
case "mime_type":
return lambda x: x.mime_type or ""
case "transfer_method":
return lambda x: x.transfer_method
Expand Down
49 changes: 46 additions & 3 deletions api/tests/unit_tests/core/workflow/nodes/test_list_operator.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,11 +2,11 @@

import pytest

from core.file import File
from core.file.models import FileTransferMethod, FileType
from core.file import File, FileTransferMethod, FileType
from core.variables import ArrayFileSegment
from core.workflow.nodes.list_operator.entities import FilterBy, FilterCondition, Limit, ListOperatorNodeData, OrderBy
from core.workflow.nodes.list_operator.node import ListOperatorNode
from core.workflow.nodes.list_operator.exc import InvalidKeyError
from core.workflow.nodes.list_operator.node import ListOperatorNode, _get_file_extract_string_func
from models.workflow import WorkflowNodeExecutionStatus


Expand Down Expand Up @@ -109,3 +109,46 @@ def test_filter_files_by_type(list_operator_node):
assert expected_file["tenant_id"] == result_file.tenant_id
assert expected_file["transfer_method"] == result_file.transfer_method
assert expected_file["related_id"] == result_file.related_id


def test_get_file_extract_string_func():
# Create a File object
file = File(
tenant_id="test_tenant",
type=FileType.DOCUMENT,
transfer_method=FileTransferMethod.LOCAL_FILE,
filename="test_file.txt",
extension=".txt",
mime_type="text/plain",
remote_url="https://example.com/test_file.txt",
related_id="test_related_id",
)

# Test each case
assert _get_file_extract_string_func(key="name")(file) == "test_file.txt"
assert _get_file_extract_string_func(key="type")(file) == "document"
assert _get_file_extract_string_func(key="extension")(file) == ".txt"
assert _get_file_extract_string_func(key="mime_type")(file) == "text/plain"
assert _get_file_extract_string_func(key="transfer_method")(file) == "local_file"
assert _get_file_extract_string_func(key="url")(file) == "https://example.com/test_file.txt"

# Test with empty values
empty_file = File(
tenant_id="test_tenant",
type=FileType.DOCUMENT,
transfer_method=FileTransferMethod.LOCAL_FILE,
filename=None,
extension=None,
mime_type=None,
remote_url=None,
related_id="test_related_id",
)

assert _get_file_extract_string_func(key="name")(empty_file) == ""
assert _get_file_extract_string_func(key="extension")(empty_file) == ""
assert _get_file_extract_string_func(key="mime_type")(empty_file) == ""
assert _get_file_extract_string_func(key="url")(empty_file) == ""

# Test invalid key
with pytest.raises(InvalidKeyError):
_get_file_extract_string_func(key="invalid_key")

0 comments on commit 07cd8d1

Please sign in to comment.