mirror of
https://github.com/Unstructured-IO/unstructured.git
synced 2025-07-25 01:50:21 +00:00

**Summary** Organize DOC tests into related groups with markers. This makes it easier to assess coverage and find tests related to particular behaviors. This is in preparation for adding tests related to DOC image extraction. No code changes, purely line-block moves. - Move module-level fixtures to the bottom. - Organize tests into related groups with markers.
387 lines
14 KiB
Python
387 lines
14 KiB
Python
import os
|
|
from tempfile import SpooledTemporaryFile
|
|
|
|
import docx
|
|
import pytest
|
|
from pytest_mock import MockFixture
|
|
|
|
from test_unstructured.unit_utils import assert_round_trips_through_JSON, example_doc_path
|
|
from unstructured.chunking.title import chunk_by_title
|
|
from unstructured.documents.elements import (
|
|
Address,
|
|
ListItem,
|
|
NarrativeText,
|
|
Table,
|
|
Text,
|
|
Title,
|
|
)
|
|
from unstructured.partition.common import convert_office_doc
|
|
from unstructured.partition.doc import partition_doc
|
|
from unstructured.partition.docx import partition_docx
|
|
|
|
|
|
def test_partition_doc_matches_partition_docx(mock_document, expected_elements, tmpdir):
|
|
docx_filename = os.path.join(tmpdir.dirname, "mock_document.docx")
|
|
doc_filename = os.path.join(tmpdir.dirname, "mock_document.doc")
|
|
mock_document.save(docx_filename)
|
|
convert_office_doc(docx_filename, tmpdir.dirname, "doc")
|
|
assert partition_doc(filename=doc_filename) == partition_docx(filename=docx_filename)
|
|
|
|
|
|
# -- document-source (file or filename) ----------------------------------------------------------
|
|
|
|
|
|
def test_partition_doc_from_filename(mock_document, expected_elements, tmpdir, capsys):
|
|
docx_filename = os.path.join(tmpdir.dirname, "mock_document.docx")
|
|
doc_filename = os.path.join(tmpdir.dirname, "mock_document.doc")
|
|
mock_document.save(docx_filename)
|
|
convert_office_doc(docx_filename, tmpdir.dirname, "doc")
|
|
elements = partition_doc(filename=doc_filename)
|
|
assert elements == expected_elements
|
|
assert elements[0].metadata.filename == "mock_document.doc"
|
|
assert elements[0].metadata.file_directory == tmpdir.dirname
|
|
assert capsys.readouterr().out == ""
|
|
assert capsys.readouterr().err == ""
|
|
|
|
|
|
def test_partition_doc_from_file_with_filter(mock_document, expected_elements, tmpdir, capsys):
|
|
docx_filename = os.path.join(tmpdir.dirname, "mock_document.docx")
|
|
doc_filename = os.path.join(tmpdir.dirname, "mock_document.doc")
|
|
mock_document.save(docx_filename)
|
|
convert_office_doc(docx_filename, tmpdir.dirname, "doc")
|
|
|
|
with open(doc_filename, "rb") as f:
|
|
elements = partition_doc(file=f, libre_office_filter="MS Word 2007 XML")
|
|
assert elements == expected_elements
|
|
assert capsys.readouterr().out == ""
|
|
assert capsys.readouterr().err == ""
|
|
for element in elements:
|
|
assert element.metadata.filename is None
|
|
|
|
|
|
def test_partition_doc_from_file_with_no_filter(mock_document, expected_elements, tmpdir, capsys):
|
|
docx_filename = os.path.join(tmpdir.dirname, "mock_document.docx")
|
|
doc_filename = os.path.join(tmpdir.dirname, "mock_document.doc")
|
|
mock_document.save(docx_filename)
|
|
convert_office_doc(docx_filename, tmpdir.dirname, "doc")
|
|
|
|
with open(doc_filename, "rb") as f:
|
|
elements = partition_doc(file=f, libre_office_filter=None)
|
|
assert elements == expected_elements
|
|
assert capsys.readouterr().out == ""
|
|
assert capsys.readouterr().err == ""
|
|
for element in elements:
|
|
assert element.metadata.filename is None
|
|
|
|
|
|
def test_partition_doc_raises_with_both_specified(mock_document, tmpdir):
|
|
docx_filename = os.path.join(tmpdir.dirname, "mock_document.docx")
|
|
doc_filename = os.path.join(tmpdir.dirname, "mock_document.doc")
|
|
mock_document.save(docx_filename)
|
|
convert_office_doc(docx_filename, tmpdir.dirname, "doc")
|
|
|
|
with open(doc_filename, "rb") as f, pytest.raises(ValueError):
|
|
partition_doc(filename=doc_filename, file=f)
|
|
|
|
|
|
def test_partition_doc_raises_with_neither():
|
|
with pytest.raises(ValueError):
|
|
partition_doc()
|
|
|
|
|
|
def test_partition_raises_with_missing_doc(mock_document, expected_elements, tmpdir):
|
|
doc_filename = os.path.join(tmpdir.dirname, "asdf.doc")
|
|
|
|
with pytest.raises(ValueError):
|
|
partition_doc(filename=doc_filename)
|
|
|
|
|
|
# -- `include_metadata` arg ----------------------------------------------------------------------
|
|
|
|
|
|
def test_partition_doc_from_filename_exclude_metadata(mock_document, tmpdir):
|
|
docx_filename = os.path.join(tmpdir.dirname, "mock_document.docx")
|
|
doc_filename = os.path.join(tmpdir.dirname, "mock_document.doc")
|
|
mock_document.save(docx_filename)
|
|
convert_office_doc(docx_filename, tmpdir.dirname, "doc")
|
|
|
|
elements = partition_doc(filename=doc_filename, include_metadata=False)
|
|
|
|
assert elements[0].metadata.filetype is None
|
|
assert elements[0].metadata.page_name is None
|
|
assert elements[0].metadata.filename is None
|
|
|
|
|
|
def test_partition_doc_from_file_exclude_metadata(mock_document, tmpdir):
|
|
docx_filename = os.path.join(tmpdir.dirname, "mock_document.docx")
|
|
doc_filename = os.path.join(tmpdir.dirname, "mock_document.doc")
|
|
mock_document.save(docx_filename)
|
|
convert_office_doc(docx_filename, tmpdir.dirname, "doc")
|
|
|
|
with open(doc_filename, "rb") as f:
|
|
elements = partition_doc(file=f, include_metadata=False)
|
|
|
|
assert elements[0].metadata.filetype is None
|
|
assert elements[0].metadata.page_name is None
|
|
assert elements[0].metadata.filename is None
|
|
|
|
|
|
# -- .metadata.filename --------------------------------------------------------------------------
|
|
|
|
|
|
def test_partition_doc_from_filename_with_metadata_filename(
|
|
mock_document,
|
|
expected_elements,
|
|
tmpdir,
|
|
):
|
|
docx_filename = os.path.join(tmpdir.dirname, "mock_document.docx")
|
|
doc_filename = os.path.join(tmpdir.dirname, "mock_document.doc")
|
|
mock_document.save(docx_filename)
|
|
convert_office_doc(docx_filename, tmpdir.dirname, "doc")
|
|
|
|
elements = partition_doc(filename=doc_filename, metadata_filename="test")
|
|
assert elements == expected_elements
|
|
assert all(element.metadata.filename == "test" for element in elements)
|
|
|
|
|
|
def test_partition_doc_from_file_with_metadata_filename(mock_document, tmpdir):
|
|
docx_filename = os.path.join(tmpdir.dirname, "mock_document.docx")
|
|
doc_filename = os.path.join(tmpdir.dirname, "mock_document.doc")
|
|
mock_document.save(docx_filename)
|
|
convert_office_doc(docx_filename, tmpdir.dirname, "doc")
|
|
|
|
with open(doc_filename, "rb") as f:
|
|
elements = partition_doc(file=f, metadata_filename="test")
|
|
for element in elements:
|
|
assert element.metadata.filename == "test"
|
|
|
|
|
|
# -- .metadata.last_modified ---------------------------------------------------------------------
|
|
|
|
|
|
def test_partition_doc_metadata_date(
|
|
mocker,
|
|
filename="example-docs/fake.doc",
|
|
):
|
|
mocked_last_modification_date = "2029-07-05T09:24:28"
|
|
|
|
mocker.patch(
|
|
"unstructured.partition.doc.get_last_modified_date",
|
|
return_value=mocked_last_modification_date,
|
|
)
|
|
|
|
elements = partition_doc(filename=filename)
|
|
|
|
assert elements[0].metadata.last_modified == mocked_last_modification_date
|
|
|
|
|
|
def test_partition_doc_metadata_date_with_custom_metadata(
|
|
mocker,
|
|
filename="example-docs/fake.doc",
|
|
):
|
|
mocked_last_modification_date = "2029-07-05T09:24:28"
|
|
expected_last_modified_date = "2020-07-05T09:24:28"
|
|
|
|
mocker.patch(
|
|
"unstructured.partition.doc.get_last_modified_date",
|
|
return_value=mocked_last_modification_date,
|
|
)
|
|
|
|
elements = partition_doc(
|
|
filename=filename,
|
|
metadata_last_modified=expected_last_modified_date,
|
|
)
|
|
|
|
assert elements[0].metadata.last_modified == expected_last_modified_date
|
|
|
|
|
|
def test_partition_doc_suppresses_modified_date_from_file_by_default(mocker: MockFixture):
|
|
modified_date_on_file = "2029-07-05T09:24:28"
|
|
mocker.patch(
|
|
"unstructured.partition.doc.get_last_modified_date_from_file",
|
|
return_value=modified_date_on_file,
|
|
)
|
|
|
|
with open(example_doc_path("fake.doc"), "rb") as f:
|
|
elements = partition_doc(file=f)
|
|
|
|
assert elements[0].metadata.last_modified is None
|
|
|
|
|
|
def test_partition_doc_pulls_modified_date_from_file_when_date_from_file_object_arg_is_True(
|
|
mocker: MockFixture,
|
|
):
|
|
modified_date_on_file = "2024-05-01T09:24:28"
|
|
mocker.patch(
|
|
"unstructured.partition.doc.get_last_modified_date_from_file",
|
|
return_value=modified_date_on_file,
|
|
)
|
|
|
|
with open(example_doc_path("fake.doc"), "rb") as f:
|
|
elements = partition_doc(file=f, date_from_file_object=True)
|
|
|
|
assert elements[0].metadata.last_modified == modified_date_on_file
|
|
|
|
|
|
def test_partition_doc_from_file_explicit_get_metadata_date(
|
|
mocker,
|
|
filename="example-docs/fake.doc",
|
|
):
|
|
mocked_last_modification_date = "2029-07-05T09:24:28"
|
|
|
|
mocker.patch(
|
|
"unstructured.partition.doc.get_last_modified_date_from_file",
|
|
return_value=mocked_last_modification_date,
|
|
)
|
|
|
|
with open(filename, "rb") as f:
|
|
elements = partition_doc(file=f, date_from_file_object=True)
|
|
|
|
assert elements[0].metadata.last_modified == mocked_last_modification_date
|
|
|
|
|
|
def test_partition_doc_from_file_without_metadata_date(
|
|
filename="example-docs/fake.doc",
|
|
):
|
|
"""Test partition_doc() with file that are not possible to get last modified date"""
|
|
with open(filename, "rb") as f:
|
|
sf = SpooledTemporaryFile()
|
|
sf.write(f.read())
|
|
sf.seek(0)
|
|
elements = partition_doc(file=sf, date_from_file_object=True)
|
|
|
|
assert elements[0].metadata.last_modified is None
|
|
|
|
|
|
def test_partition_doc_from_file_metadata_date_with_custom_metadata(
|
|
mocker,
|
|
filename="example-docs/fake.doc",
|
|
):
|
|
mocked_last_modification_date = "2029-07-05T09:24:28"
|
|
expected_last_modified_date = "2020-07-05T09:24:28"
|
|
|
|
mocker.patch(
|
|
"unstructured.partition.doc.get_last_modified_date_from_file",
|
|
return_value=mocked_last_modification_date,
|
|
)
|
|
with open(filename, "rb") as f:
|
|
elements = partition_doc(
|
|
file=f, metadata_last_modified=expected_last_modified_date, date_from_file_object=True
|
|
)
|
|
|
|
assert elements[0].metadata.last_modified == expected_last_modified_date
|
|
|
|
|
|
# -- language-recognition metadata ---------------------------------------------------------------
|
|
|
|
|
|
def test_partition_doc_element_metadata_has_languages():
|
|
filename = "example-docs/fake-doc-emphasized-text.doc"
|
|
elements = partition_doc(filename=filename)
|
|
assert elements[0].metadata.languages == ["eng"]
|
|
|
|
|
|
def test_partition_doc_respects_detect_language_per_element():
|
|
filename = "example-docs/language-docs/eng_spa_mult.doc"
|
|
elements = partition_doc(filename=filename, detect_language_per_element=True)
|
|
langs = [element.metadata.languages for element in elements]
|
|
assert langs == [["eng"], ["spa", "eng"], ["eng"], ["eng"], ["spa"]]
|
|
|
|
|
|
# -- miscellaneous -------------------------------------------------------------------------------
|
|
|
|
|
|
def test_partition_doc_grabs_emphasized_texts():
|
|
expected_emphasized_text_contents = ["bold", "italic", "bold-italic", "bold-italic"]
|
|
expected_emphasized_text_tags = ["b", "i", "b", "i"]
|
|
|
|
elements = partition_doc("example-docs/fake-doc-emphasized-text.doc")
|
|
|
|
assert isinstance(elements[0], Table)
|
|
assert elements[0].metadata.emphasized_text_contents == expected_emphasized_text_contents
|
|
assert elements[0].metadata.emphasized_text_tags == expected_emphasized_text_tags
|
|
|
|
assert elements[1] == NarrativeText("I am a bold italic bold-italic text.")
|
|
assert elements[1].metadata.emphasized_text_contents == expected_emphasized_text_contents
|
|
assert elements[1].metadata.emphasized_text_tags == expected_emphasized_text_tags
|
|
|
|
assert elements[2] == NarrativeText("I am a normal text.")
|
|
assert elements[2].metadata.emphasized_text_contents is None
|
|
assert elements[2].metadata.emphasized_text_tags is None
|
|
|
|
|
|
def test_partition_doc_with_json(mock_document, tmpdir):
|
|
docx_filename = os.path.join(tmpdir.dirname, "mock_document.docx")
|
|
doc_filename = os.path.join(tmpdir.dirname, "mock_document.doc")
|
|
mock_document.save(docx_filename)
|
|
convert_office_doc(docx_filename, tmpdir.dirname, "doc")
|
|
|
|
elements = partition_doc(filename=doc_filename)
|
|
assert_round_trips_through_JSON(elements)
|
|
|
|
|
|
def test_add_chunking_strategy_on_partition_doc(filename="example-docs/fake.doc"):
|
|
chunk_elements = partition_doc(filename, chunking_strategy="by_title")
|
|
elements = partition_doc(filename)
|
|
chunks = chunk_by_title(elements)
|
|
assert chunk_elements != elements
|
|
assert chunk_elements == chunks
|
|
|
|
|
|
def test_partition_doc_for_deterministic_and_unique_ids():
|
|
ids = [element.id for element in partition_doc("example-docs/duplicate-paragraphs.doc")]
|
|
|
|
assert ids == [
|
|
"ade273c622c48d67a7be7b3816d5b4d8",
|
|
"7d0b32fdf169f9578723486cb4bc1235",
|
|
"1feb6e8e9c1662cfaef75907aeeb0900",
|
|
"aa2a8ac10143b12f0fe2087837ea11d2",
|
|
"da31ba7ed3919067d2c6572dc1617271",
|
|
"1914359c179a160df921b769acf8c353",
|
|
"f9d0d379fc791bae487b7a45f65caa50",
|
|
]
|
|
|
|
|
|
# == module-level fixtures =======================================================================
|
|
|
|
|
|
@pytest.fixture()
|
|
def expected_elements():
|
|
return [
|
|
Title("These are a few of my favorite things:"),
|
|
ListItem("Parrots"),
|
|
ListItem("Hockey"),
|
|
Title("Analysis"),
|
|
NarrativeText("This is my first thought. This is my second thought."),
|
|
NarrativeText("This is my third thought."),
|
|
Text("2023"),
|
|
Address("DOYLESTOWN, PA 18901"),
|
|
]
|
|
|
|
|
|
@pytest.fixture()
|
|
def mock_document():
|
|
document = docx.Document()
|
|
|
|
document.add_paragraph("These are a few of my favorite things:", style="Heading 1")
|
|
# NOTE(robinson) - this should get picked up as a list item due to the •
|
|
document.add_paragraph("• Parrots", style="Normal")
|
|
# NOTE(robinson) - this should get dropped because it's empty
|
|
document.add_paragraph("• ", style="Normal")
|
|
document.add_paragraph("Hockey", style="List Bullet")
|
|
# NOTE(robinson) - this should get dropped because it's empty
|
|
document.add_paragraph("", style="List Bullet")
|
|
# NOTE(robinson) - this should get picked up as a title
|
|
document.add_paragraph("Analysis", style="Normal")
|
|
# NOTE(robinson) - this should get dropped because it is empty
|
|
document.add_paragraph("", style="Normal")
|
|
# NOTE(robinson) - this should get picked up as a narrative text
|
|
document.add_paragraph("This is my first thought. This is my second thought.", style="Normal")
|
|
document.add_paragraph("This is my third thought.", style="Body Text")
|
|
# NOTE(robinson) - this should just be regular text
|
|
document.add_paragraph("2023")
|
|
# NOTE(robinson) - this should be an address
|
|
document.add_paragraph("DOYLESTOWN, PA 18901")
|
|
|
|
return document
|