mirror of
https://git.mirrors.martin98.com/https://github.com/langgenius/dify.git
synced 2025-08-12 16:59:04 +08:00
Co-authored-by: wangheyang <wangheyang@corp.netease.com> Co-authored-by: crazywoola <427733928@qq.com>
This commit is contained in:
parent
38e48c0c40
commit
b9b5d43dc6
@ -366,7 +366,7 @@ def _extract_text_from_excel(file_content: bytes) -> str:
|
|||||||
df = excel_file.parse(sheet_name=sheet_name)
|
df = excel_file.parse(sheet_name=sheet_name)
|
||||||
df.dropna(how="all", inplace=True)
|
df.dropna(how="all", inplace=True)
|
||||||
# Create Markdown table two times to separate tables with a newline
|
# Create Markdown table two times to separate tables with a newline
|
||||||
markdown_table += df.to_markdown(index=False) + "\n\n"
|
markdown_table += df.to_markdown(index=False, floatfmt="") + "\n\n"
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
continue
|
continue
|
||||||
return markdown_table
|
return markdown_table
|
||||||
|
@ -10,6 +10,7 @@ from core.workflow.entities.node_entities import NodeRunResult
|
|||||||
from core.workflow.nodes.document_extractor import DocumentExtractorNode, DocumentExtractorNodeData
|
from core.workflow.nodes.document_extractor import DocumentExtractorNode, DocumentExtractorNodeData
|
||||||
from core.workflow.nodes.document_extractor.node import (
|
from core.workflow.nodes.document_extractor.node import (
|
||||||
_extract_text_from_docx,
|
_extract_text_from_docx,
|
||||||
|
_extract_text_from_excel,
|
||||||
_extract_text_from_pdf,
|
_extract_text_from_pdf,
|
||||||
_extract_text_from_plain_text,
|
_extract_text_from_plain_text,
|
||||||
)
|
)
|
||||||
@ -182,3 +183,181 @@ def test_extract_text_from_docx(mock_document):
|
|||||||
|
|
||||||
def test_node_type(document_extractor_node):
|
def test_node_type(document_extractor_node):
|
||||||
assert document_extractor_node._node_type == NodeType.DOCUMENT_EXTRACTOR
|
assert document_extractor_node._node_type == NodeType.DOCUMENT_EXTRACTOR
|
||||||
|
|
||||||
|
|
||||||
|
@patch("pandas.ExcelFile")
|
||||||
|
def test_extract_text_from_excel_single_sheet(mock_excel_file):
|
||||||
|
"""Test extracting text from Excel file with single sheet."""
|
||||||
|
# Mock DataFrame
|
||||||
|
mock_df = Mock()
|
||||||
|
mock_df.dropna = Mock()
|
||||||
|
mock_df.to_markdown.return_value = "| Name | Age |\n|------|-----|\n| John | 25 |"
|
||||||
|
|
||||||
|
# Mock ExcelFile
|
||||||
|
mock_excel_instance = Mock()
|
||||||
|
mock_excel_instance.sheet_names = ["Sheet1"]
|
||||||
|
mock_excel_instance.parse.return_value = mock_df
|
||||||
|
mock_excel_file.return_value = mock_excel_instance
|
||||||
|
|
||||||
|
file_content = b"fake_excel_content"
|
||||||
|
result = _extract_text_from_excel(file_content)
|
||||||
|
|
||||||
|
expected = "| Name | Age |\n|------|-----|\n| John | 25 |\n\n"
|
||||||
|
assert result == expected
|
||||||
|
mock_excel_file.assert_called_once()
|
||||||
|
mock_df.dropna.assert_called_once_with(how="all", inplace=True)
|
||||||
|
mock_df.to_markdown.assert_called_once_with(index=False, floatfmt="")
|
||||||
|
|
||||||
|
|
||||||
|
@patch("pandas.ExcelFile")
|
||||||
|
def test_extract_text_from_excel_multiple_sheets(mock_excel_file):
|
||||||
|
"""Test extracting text from Excel file with multiple sheets."""
|
||||||
|
# Mock DataFrames for different sheets
|
||||||
|
mock_df1 = Mock()
|
||||||
|
mock_df1.dropna = Mock()
|
||||||
|
mock_df1.to_markdown.return_value = "| Product | Price |\n|---------|-------|\n| Apple | 1.50 |"
|
||||||
|
|
||||||
|
mock_df2 = Mock()
|
||||||
|
mock_df2.dropna = Mock()
|
||||||
|
mock_df2.to_markdown.return_value = "| City | Population |\n|------|------------|\n| NYC | 8000000 |"
|
||||||
|
|
||||||
|
# Mock ExcelFile
|
||||||
|
mock_excel_instance = Mock()
|
||||||
|
mock_excel_instance.sheet_names = ["Products", "Cities"]
|
||||||
|
mock_excel_instance.parse.side_effect = [mock_df1, mock_df2]
|
||||||
|
mock_excel_file.return_value = mock_excel_instance
|
||||||
|
|
||||||
|
file_content = b"fake_excel_content_multiple_sheets"
|
||||||
|
result = _extract_text_from_excel(file_content)
|
||||||
|
|
||||||
|
expected = (
|
||||||
|
"| Product | Price |\n|---------|-------|\n| Apple | 1.50 |\n\n"
|
||||||
|
"| City | Population |\n|------|------------|\n| NYC | 8000000 |\n\n"
|
||||||
|
)
|
||||||
|
assert result == expected
|
||||||
|
assert mock_excel_instance.parse.call_count == 2
|
||||||
|
|
||||||
|
|
||||||
|
@patch("pandas.ExcelFile")
|
||||||
|
def test_extract_text_from_excel_empty_sheets(mock_excel_file):
|
||||||
|
"""Test extracting text from Excel file with empty sheets."""
|
||||||
|
# Mock empty DataFrame
|
||||||
|
mock_df = Mock()
|
||||||
|
mock_df.dropna = Mock()
|
||||||
|
mock_df.to_markdown.return_value = ""
|
||||||
|
|
||||||
|
# Mock ExcelFile
|
||||||
|
mock_excel_instance = Mock()
|
||||||
|
mock_excel_instance.sheet_names = ["EmptySheet"]
|
||||||
|
mock_excel_instance.parse.return_value = mock_df
|
||||||
|
mock_excel_file.return_value = mock_excel_instance
|
||||||
|
|
||||||
|
file_content = b"fake_excel_empty_content"
|
||||||
|
result = _extract_text_from_excel(file_content)
|
||||||
|
|
||||||
|
expected = "\n\n"
|
||||||
|
assert result == expected
|
||||||
|
|
||||||
|
|
||||||
|
@patch("pandas.ExcelFile")
|
||||||
|
def test_extract_text_from_excel_sheet_parse_error(mock_excel_file):
|
||||||
|
"""Test handling of sheet parsing errors - should continue with other sheets."""
|
||||||
|
# Mock DataFrames - one successful, one that raises exception
|
||||||
|
mock_df_success = Mock()
|
||||||
|
mock_df_success.dropna = Mock()
|
||||||
|
mock_df_success.to_markdown.return_value = "| Data | Value |\n|------|-------|\n| Test | 123 |"
|
||||||
|
|
||||||
|
# Mock ExcelFile
|
||||||
|
mock_excel_instance = Mock()
|
||||||
|
mock_excel_instance.sheet_names = ["GoodSheet", "BadSheet"]
|
||||||
|
mock_excel_instance.parse.side_effect = [mock_df_success, Exception("Parse error")]
|
||||||
|
mock_excel_file.return_value = mock_excel_instance
|
||||||
|
|
||||||
|
file_content = b"fake_excel_mixed_content"
|
||||||
|
result = _extract_text_from_excel(file_content)
|
||||||
|
|
||||||
|
expected = "| Data | Value |\n|------|-------|\n| Test | 123 |\n\n"
|
||||||
|
assert result == expected
|
||||||
|
|
||||||
|
|
||||||
|
@patch("pandas.ExcelFile")
|
||||||
|
def test_extract_text_from_excel_file_error(mock_excel_file):
|
||||||
|
"""Test handling of Excel file reading errors."""
|
||||||
|
mock_excel_file.side_effect = Exception("Invalid Excel file")
|
||||||
|
|
||||||
|
file_content = b"invalid_excel_content"
|
||||||
|
|
||||||
|
with pytest.raises(Exception) as exc_info:
|
||||||
|
_extract_text_from_excel(file_content)
|
||||||
|
|
||||||
|
# Note: The function should raise TextExtractionError, but since it's not imported in the test,
|
||||||
|
# we check for the general Exception pattern
|
||||||
|
assert "Failed to extract text from Excel file" in str(exc_info.value)
|
||||||
|
|
||||||
|
|
||||||
|
@patch("pandas.ExcelFile")
|
||||||
|
def test_extract_text_from_excel_io_bytesio_usage(mock_excel_file):
|
||||||
|
"""Test that BytesIO is properly used with the file content."""
|
||||||
|
import io
|
||||||
|
|
||||||
|
# Mock DataFrame
|
||||||
|
mock_df = Mock()
|
||||||
|
mock_df.dropna = Mock()
|
||||||
|
mock_df.to_markdown.return_value = "| Test | Data |\n|------|------|\n| 1 | A |"
|
||||||
|
|
||||||
|
# Mock ExcelFile
|
||||||
|
mock_excel_instance = Mock()
|
||||||
|
mock_excel_instance.sheet_names = ["TestSheet"]
|
||||||
|
mock_excel_instance.parse.return_value = mock_df
|
||||||
|
mock_excel_file.return_value = mock_excel_instance
|
||||||
|
|
||||||
|
file_content = b"test_excel_bytes"
|
||||||
|
result = _extract_text_from_excel(file_content)
|
||||||
|
|
||||||
|
# Verify that ExcelFile was called with a BytesIO object
|
||||||
|
mock_excel_file.assert_called_once()
|
||||||
|
call_args = mock_excel_file.call_args[0][0]
|
||||||
|
assert isinstance(call_args, io.BytesIO)
|
||||||
|
|
||||||
|
expected = "| Test | Data |\n|------|------|\n| 1 | A |\n\n"
|
||||||
|
assert result == expected
|
||||||
|
|
||||||
|
|
||||||
|
@patch("pandas.ExcelFile")
|
||||||
|
def test_extract_text_from_excel_all_sheets_fail(mock_excel_file):
|
||||||
|
"""Test when all sheets fail to parse - should return empty string."""
|
||||||
|
# Mock ExcelFile
|
||||||
|
mock_excel_instance = Mock()
|
||||||
|
mock_excel_instance.sheet_names = ["BadSheet1", "BadSheet2"]
|
||||||
|
mock_excel_instance.parse.side_effect = [Exception("Error 1"), Exception("Error 2")]
|
||||||
|
mock_excel_file.return_value = mock_excel_instance
|
||||||
|
|
||||||
|
file_content = b"fake_excel_all_bad_sheets"
|
||||||
|
result = _extract_text_from_excel(file_content)
|
||||||
|
|
||||||
|
# Should return empty string when all sheets fail
|
||||||
|
assert result == ""
|
||||||
|
|
||||||
|
|
||||||
|
@patch("pandas.ExcelFile")
|
||||||
|
def test_extract_text_from_excel_markdown_formatting(mock_excel_file):
|
||||||
|
"""Test that markdown formatting parameters are correctly applied."""
|
||||||
|
# Mock DataFrame
|
||||||
|
mock_df = Mock()
|
||||||
|
mock_df.dropna = Mock()
|
||||||
|
mock_df.to_markdown.return_value = "| Float | Int |\n|-------|-----|\n| 123456.78 | 42 |"
|
||||||
|
|
||||||
|
# Mock ExcelFile
|
||||||
|
mock_excel_instance = Mock()
|
||||||
|
mock_excel_instance.sheet_names = ["NumberSheet"]
|
||||||
|
mock_excel_instance.parse.return_value = mock_df
|
||||||
|
mock_excel_file.return_value = mock_excel_instance
|
||||||
|
|
||||||
|
file_content = b"fake_excel_numbers"
|
||||||
|
result = _extract_text_from_excel(file_content)
|
||||||
|
|
||||||
|
# Verify to_markdown was called with correct parameters
|
||||||
|
mock_df.to_markdown.assert_called_once_with(index=False, floatfmt="")
|
||||||
|
|
||||||
|
expected = "| Float | Int |\n|-------|-----|\n| 123456.78 | 42 |\n\n"
|
||||||
|
assert result == expected
|
||||||
|
Loading…
x
Reference in New Issue
Block a user