unstructured/test_unstructured/partition/test_xml_partition.py
John e9fdbb0943
feat: add include_metadata across all partition functions (#853)
* add include_metadata kwarg and tests to parsers

add exclude_metadata to docx

add test for doc to exclude metadata

add include_metadata kwarg to email

add include_metadata kwarg to epub

add include_metadata kwarg to json

add exclude_metadata tests to md

add include_metadata kwarg and tests for msg parse

add include_metadata kwarg and tests for odt parse

add include_metadata kwarg and tests for org parse

add include_metadata kwarg and tests for ppt and pptx parse

add include_metadata kwarg and tests for rst parse

add include_metadata kwarg and tests for rtf parse

add include_metadata tests for text parse

add include_metadata tests for tsv parse

add include_metadata tests for xlsx parse

add include_metadata tests for xml parse

* WIP add include_metadata to partition_pdf

* add include_metadata tests to partition_pdf

* make tidy/check

* update changelog and version

* change test asserts and move docstring logic to process_metadata

* make tidy

* fix tests asserts

* linting, linting, linting

* sync versions

* skip api call test not on main

---------

Co-authored-by: Matt Robinson <mrobinson@unstructured.io>
Co-authored-by: Matt Robinson <mrobinson@unstructuredai.io>
2023-06-30 10:44:46 -04:00

142 lines
4.7 KiB
Python

import os
import pathlib
import pytest
from unstructured.partition.xml import partition_xml
DIRECTORY = pathlib.Path(__file__).parent.resolve()
@pytest.mark.parametrize(
"filename",
["factbook.xml", "factbook-utf-16.xml"],
)
def test_partition_xml_from_filename(filename):
file_path = os.path.join(DIRECTORY, "..", "..", "example-docs", filename)
elements = partition_xml(filename=file_path, xml_keep_tags=False)
assert elements[0].text == "United States"
assert elements[0].metadata.filename == filename
@pytest.mark.parametrize(
"filename",
["factbook.xml", "factbook-utf-16.xml"],
)
def test_partition_xml_from_file(filename):
file_path = os.path.join(DIRECTORY, "..", "..", "example-docs", filename)
with open(file_path) as f:
elements = partition_xml(file=f, xml_keep_tags=False, metadata_filename=file_path)
assert elements[0].text == "United States"
assert elements[0].metadata.filename == filename
@pytest.mark.parametrize(
"filename",
["factbook.xml", "factbook-utf-16.xml"],
)
def test_partition_xml_from_file_rb(filename):
file_path = os.path.join(DIRECTORY, "..", "..", "example-docs", filename)
with open(file_path, "rb") as f:
elements = partition_xml(file=f, xml_keep_tags=False, metadata_filename=file_path)
assert elements[0].text == "United States"
assert elements[0].metadata.filename == filename
@pytest.mark.parametrize(
"filename",
["factbook.xml", "factbook-utf-16.xml"],
)
def test_partition_xml_from_filename_with_tags_default_encoding(filename):
file_path = os.path.join(DIRECTORY, "..", "..", "example-docs", filename)
elements = partition_xml(filename=file_path, xml_keep_tags=True)
assert elements[5].text == "<name>United States</name>"
assert elements[5].metadata.filename == filename
@pytest.mark.parametrize(
("filename", "encoding", "error"),
[("factbook-utf-16.xml", "utf-8", UnicodeDecodeError)],
)
def test_partition_xml_from_filename_with_tags_raises_encoding_error(filename, encoding, error):
with pytest.raises(error):
file_path = os.path.join(DIRECTORY, "..", "..", "example-docs", filename)
partition_xml(filename=file_path, xml_keep_tags=True, encoding=encoding)
@pytest.mark.parametrize(
"filename",
["factbook.xml", "factbook-utf-16.xml"],
)
def test_partition_xml_from_file_with_tags_default_encoding(filename):
file_path = os.path.join(DIRECTORY, "..", "..", "example-docs", filename)
with open(file_path) as f:
elements = partition_xml(file=f, xml_keep_tags=True, metadata_filename=file_path)
assert elements[5].text == "<name>United States</name>"
assert elements[5].metadata.filename == filename
@pytest.mark.parametrize(
"filename",
["factbook.xml", "factbook-utf-16.xml"],
)
def test_partition_xml_from_file_rb_with_tags_default_encoding(filename):
file_path = os.path.join(DIRECTORY, "..", "..", "example-docs", filename)
with open(file_path, "rb") as f:
elements = partition_xml(file=f, xml_keep_tags=True, metadata_filename=file_path)
assert elements[5].text == "<name>United States</name>"
assert elements[5].metadata.filename == filename
@pytest.mark.parametrize(
("filename", "encoding", "error"),
[("factbook-utf-16.xml", "utf-8", UnicodeDecodeError)],
)
def test_partition_xml_from_file_rb_with_tags_raises_encoding_error(filename, encoding, error):
file_path = os.path.join(DIRECTORY, "..", "..", "example-docs", filename)
with pytest.raises(error), open(file_path, "rb") as f:
partition_xml(
file=f,
xml_keep_tags=True,
metadata_filename=file_path,
encoding=encoding,
)
@pytest.mark.parametrize(
"filename",
["factbook.xml", "factbook-utf-16.xml"],
)
def test_partition_xml_from_filename_exclude_metadata(filename):
file_path = os.path.join(DIRECTORY, "..", "..", "example-docs", filename)
elements = partition_xml(filename=file_path, xml_keep_tags=False, include_metadata=False)
assert elements[0].text == "United States"
for i in range(len(elements)):
assert elements[i].metadata.to_dict() == {}
@pytest.mark.parametrize(
"filename",
["factbook.xml", "factbook-utf-16.xml"],
)
def test_partition_xml_from_file_exclude_metadata(filename):
file_path = os.path.join(DIRECTORY, "..", "..", "example-docs", filename)
with open(file_path) as f:
elements = partition_xml(
file=f,
xml_keep_tags=False,
metadata_filename=file_path,
include_metadata=False,
)
assert elements[0].text == "United States"
for i in range(len(elements)):
assert elements[i].metadata.to_dict() == {}