hz/ebiz-dify-ai
1
0
Fork 0
mirror of http://112.124.100.131/huang.ze/ebiz-dify-ai.git synced 2025-12-10 03:16:51 +08:00
Code Issues Packages Projects Releases Wiki Activity

chore: make doc extractor node also can extract text by file extension (#9543)

Browse Source
This commit is contained in:
非法操作
2024-10-24 22:54:48 +08:00
committed by GitHub
parent 88dec6ef2b
commit 2346b0ab99
2 changed files with 48 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

19
api/tests/unit_tests/core/workflow/nodes/test_document_extractor_node.py
View File

@@ -63,17 +63,24 @@ def test_run_invalid_variable_type(document_extractor_node, mock_graph_runtime_s
@pytest.mark.parametrize(
("mime_type", "file_content", "expected_text", "transfer_method"),
("mime_type", "file_content", "expected_text", "transfer_method", "extension"),
[
("text/plain", b"Hello, world!", ["Hello, world!"], FileTransferMethod.LOCAL_FILE),
("application/pdf", b"%PDF-1.5\n%Test PDF content", ["Mocked PDF content"], FileTransferMethod.LOCAL_FILE),
("text/plain", b"Hello, world!", ["Hello, world!"], FileTransferMethod.LOCAL_FILE, ".txt"),
(
"application/pdf",
b"%PDF-1.5\n%Test PDF content",
["Mocked PDF content"],
FileTransferMethod.LOCAL_FILE,
".pdf",
),
(
"application/vnd.openxmlformats-officedocument.wordprocessingml.document",
b"PK\x03\x04",
["Mocked DOCX content"],
FileTransferMethod.LOCAL_FILE,
FileTransferMethod.REMOTE_URL,
"",
),
("text/plain", b"Remote content", ["Remote content"], FileTransferMethod.REMOTE_URL),
("text/plain", b"Remote content", ["Remote content"], FileTransferMethod.REMOTE_URL, None),
],
)
def test_run_extract_text(
@@ -83,6 +90,7 @@ def test_run_extract_text(
file_content,
expected_text,
transfer_method,
extension,
monkeypatch,
):
document_extractor_node.graph_runtime_state = mock_graph_runtime_state
@@ -92,6 +100,7 @@ def test_run_extract_text(
mock_file.transfer_method = transfer_method
mock_file.related_id = "test_file_id" if transfer_method == FileTransferMethod.LOCAL_FILE else None
mock_file.remote_url = "https://example.com/file.txt" if transfer_method == FileTransferMethod.REMOTE_URL else None
mock_file.extension = extension
mock_array_file_segment = Mock(spec=ArrayFileSegment)
mock_array_file_segment.value = [mock_file]