From ea38e484b81b4104736d2484aaa1d0d7ea99627e Mon Sep 17 00:00:00 2001 From: Asuka Minato Date: Wed, 30 Sep 2026 03:44:01 +0000 Subject: [PATCH] test: use real document extractor file fixtures (#43156) --- .../nodes/test_document_extractor_node.py | 85 +++++++++---------- 1 file changed, 40 insertions(+), 45 deletions(-) diff --git a/api/tests/unit_tests/core/workflow/nodes/test_document_extractor_node.py b/api/tests/unit_tests/core/workflow/nodes/test_document_extractor_node.py index ee094732b65..d8bc859dae6 100644 --- a/api/tests/unit_tests/core/workflow/nodes/test_document_extractor_node.py +++ b/api/tests/unit_tests/core/workflow/nodes/test_document_extractor_node.py @@ -3,12 +3,12 @@ from unittest.mock import Mock, patch import pandas as pd import pytest -from docx.oxml.text.paragraph import CT_P +from docx import Document from core.app.entities.app_invoke_entities import InvokeFrom, UserFrom from graphon.entities import GraphInitParams from graphon.enums import BuiltinNodeTypes, WorkflowNodeExecutionStatus -from graphon.file import File, FileTransferMethod +from graphon.file import File, FileTransferMethod, FileType from graphon.node_events import NodeRunResult from graphon.nodes.document_extractor import DocumentExtractorNode, DocumentExtractorNodeData from graphon.nodes.document_extractor.exc import TextExtractionError, UnsupportedFileTypeError @@ -26,6 +26,17 @@ from graphon.variables.variables import StringVariable from tests.workflow_test_utils import build_test_graph_init_params +def make_file(extension: str | None, mime_type: str | None) -> File: + """Use None when the file has no extension or MIME metadata, respectively.""" + return File( + file_type=FileType.DOCUMENT, + transfer_method=FileTransferMethod.LOCAL_FILE, + reference="test_file_id", + extension=extension, + mime_type=mime_type, + ) + + @pytest.fixture def graph_init_params() -> GraphInitParams: return build_test_graph_init_params( @@ -107,9 +118,8 @@ def test_run_none_only_file_list_returns_succeeded(document_extractor_node, mock """A file list containing only None (e.g., [None]) should be filtered to [] and succeed.""" document_extractor_node.graph_runtime_state = mock_graph_runtime_state - # Use a Mock to bypass type validation for None entries in the list - afs = Mock(spec=ArrayFileSegment) - afs.value = [None] + # Preserve the legacy null-entry payload handled by the node filtering branch. + afs = ArrayFileSegment.model_construct(value=[None]) mock_graph_runtime_state.variable_pool.get.return_value = afs result = document_extractor_node._run() @@ -165,17 +175,17 @@ def test_run_extract_text( ): document_extractor_node.graph_runtime_state = mock_graph_runtime_state - mock_file = Mock(spec=File) - mock_file.mime_type = mime_type - 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 + file = File( + file_type=FileType.DOCUMENT, + mime_type=mime_type, + transfer_method=transfer_method, + reference="test_file_id" if transfer_method == FileTransferMethod.LOCAL_FILE else None, + remote_url="https://example.com/file.txt" if transfer_method == FileTransferMethod.REMOTE_URL else None, + extension=extension, + ) + array_file_segment = ArrayFileSegment(value=[file]) - mock_array_file_segment = Mock(spec=ArrayFileSegment) - mock_array_file_segment.value = [mock_file] - - mock_graph_runtime_state.variable_pool.get.return_value = mock_array_file_segment + mock_graph_runtime_state.variable_pool.get.return_value = array_file_segment mock_download = Mock(return_value=file_content) @@ -217,7 +227,7 @@ def test_run_extract_text( if transfer_method == FileTransferMethod.REMOTE_URL: document_extractor_node._http_client.get.assert_called_once_with("https://example.com/file.txt") elif transfer_method == FileTransferMethod.LOCAL_FILE: - mock_download.assert_called_once_with(mock_file) + mock_download.assert_called_once_with(file) def test_extract_text_from_plain_text(): @@ -247,20 +257,14 @@ def test_extract_text_from_pdf(mock_pdf_document): assert text == "PDF content" -@patch("docx.Document") -def test_extract_text_from_docx(mock_document): - mock_paragraph1 = Mock() - mock_paragraph1.text = "Paragraph 1" - mock_paragraph2 = Mock() - mock_paragraph2.text = "Paragraph 2" - mock_document.return_value.paragraphs = [mock_paragraph1, mock_paragraph2] - mock_ct_p1 = Mock(spec=CT_P) - mock_ct_p1.text = "Paragraph 1" - mock_ct_p2 = Mock(spec=CT_P) - mock_ct_p2.text = "Paragraph 2" - mock_element = Mock(body=[mock_ct_p1, mock_ct_p2]) - mock_document.return_value.element = mock_element - text = _extract_text_from_docx(b"PK\x03\x04") +def test_extract_text_from_docx(): + document = Document() + document.add_paragraph("Paragraph 1") + document.add_paragraph("Paragraph 2") + content = io.BytesIO() + document.save(content) + + text = _extract_text_from_docx(content.getvalue()) assert text == "Paragraph 1\nParagraph 2" @@ -454,9 +458,7 @@ def test_extract_text_from_file_routes_odt_inputs_to_graphon_odt_extractor( mime_type, route_label, ): - file = Mock(spec=File) - file.extension = extension - file.mime_type = mime_type + file = make_file(extension, mime_type) def fake_partition(file_content, *, suffix, unstructured_api_config, load_local_partition, render_element): assert file_content == b"odt content" @@ -491,9 +493,7 @@ def test_extract_text_from_file_routes_odt_inputs_to_graphon_odt_extractor( ], ) def test_extract_text_from_file_routes_excel_inputs(document_extractor_node, extension, mime_type): - file = Mock(spec=File) - file.extension = extension - file.mime_type = mime_type + file = make_file(extension, mime_type) with patch( "graphon.nodes.document_extractor.node._download_file_content", @@ -534,9 +534,7 @@ def test_extract_text_from_file_routes_excel_inputs(document_extractor_node, ext def test_extract_text_from_file_rejects_missing_extension_and_mime_type(document_extractor_node): - file = Mock(spec=File) - file.extension = None - file.mime_type = None + file = make_file(None, None) with patch( "graphon.nodes.document_extractor.node._download_file_content", @@ -552,8 +550,7 @@ def test_extract_text_from_file_rejects_missing_extension_and_mime_type(document def test_run_list_file_extraction_error_returns_failed(document_extractor_node, mock_graph_runtime_state): document_extractor_node.graph_runtime_state = mock_graph_runtime_state - file_list = Mock(spec=ArrayFileSegment) - file_list.value = [Mock(spec=File)] + file_list = ArrayFileSegment(value=[make_file(".txt", "text/plain")]) mock_graph_runtime_state.variable_pool.get.return_value = file_list with patch( @@ -568,8 +565,7 @@ def test_run_list_file_extraction_error_returns_failed(document_extractor_node, def test_run_single_file_segment_extraction_error_returns_failed(document_extractor_node, mock_graph_runtime_state): document_extractor_node.graph_runtime_state = mock_graph_runtime_state - file_segment = Mock(spec=FileSegment) - file_segment.value = Mock(spec=File) + file_segment = FileSegment(value=make_file(".txt", "text/plain")) mock_graph_runtime_state.variable_pool.get.return_value = file_segment with patch( @@ -584,8 +580,7 @@ def test_run_single_file_segment_extraction_error_returns_failed(document_extrac def test_run_single_file_segment_returns_string_output(document_extractor_node, mock_graph_runtime_state): document_extractor_node.graph_runtime_state = mock_graph_runtime_state - file_segment = Mock(spec=FileSegment) - file_segment.value = Mock(spec=File) + file_segment = FileSegment(value=make_file(".txt", "text/plain")) mock_graph_runtime_state.variable_pool.get.return_value = file_segment with patch(