Steve Canny db186dc23b
rfctr(doc): organize test_doc.py (#3017)
**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.
2024-05-14 20:57:31 +00:00

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