unstructured

1
# pyright: reportPrivateUsage=false
2

3
import os
4
from pathlib import Path
5

6
import pytest
7
from lxml import etree
8

9
from unstructured.documents.xml import XMLDocument
10

11
FILEPATH = Path(__file__).absolute().parent
12

13

14
@pytest.fixture()
15
def sample_document():
16
    return """"<SEC-DOCUMENT>
17
    <TYPE>10-K
18
    <COMPANY>Proctor & Gamble
19
</SEC-DOCUMENT>"""
20

21

22
def test_read_xml(sample_document, tmpdir):
23
    filename = os.path.join(tmpdir.dirname, "sample-document.xml")
24
    with open(filename, "w") as f:
25
        f.write(sample_document)
26

27
    xml_document = XMLDocument.from_file(filename=filename)
28
    document_tree = xml_document.document_tree
29
    type_tag = document_tree.find(".//type")
30
    assert type_tag.text.strip() == "10-K"
31

32
    # Test to make sure the & character is retained
33
    company_tag = document_tree.find(".//company")
34
    assert company_tag.text.strip() == "Proctor & Gamble"
35

36

37
def test_xml_read_raises():
38
    xml_document = XMLDocument()
39
    with pytest.raises(NotImplementedError):
40
        xml_document._parse_pages_from_element_tree()
41

42

43
def test_from_string(sample_document):
44
    xml_document = XMLDocument.from_string(sample_document)
45
    type_tag = xml_document.document_tree.find(".//type")
46
    assert type_tag.text.strip() == "10-K"
47

48

49
def test_from_string_with_pre_tag():
50
    sample_document = """
51
    <pre>
52
    <SEC-DOCUMENT>
53
    <TYPE>10-K
54
    <COMPANY>Proctor & Gamble
55
    </SEC-DOCUMENT>
56
    </pre>
57
    """
58
    xml_document = XMLDocument.from_string(sample_document)
59
    type_tag = xml_document.document_tree.find(".//type")
60
    assert type_tag.text.strip() == "10-K"
61

62

63
def test_read_with_stylesheet():
64
    filename = os.path.join(FILEPATH, "..", "..", "example-docs", "factbook.xml")
65
    stylesheet = os.path.join(FILEPATH, "..", "..", "example-docs", "unsupported", "factbook.xsl")
66

67
    xml_document = XMLDocument.from_file(filename=filename, stylesheet=stylesheet)
68
    doc_tree = xml_document.document_tree
69
    # NOTE(robinson) - The table heading row plus one row for each of the four data items
70
    assert int(doc_tree.xpath("count(//tr)")) == 5
71
    # NOTE(robinson) - Four data elements x four attributes for each
72
    assert int(doc_tree.xpath("count(//td)")) == 16
73

74

75
def test_read_with_stylesheet_warns_with_html_parser(caplog):
76
    filename = os.path.join(FILEPATH, "..", "..", "example-docs", "factbook.xml")
77
    stylesheet = os.path.join(FILEPATH, "..", "..", "example-docs", "unsupported", "factbook.xsl")
78

79
    XMLDocument.from_file(filename=filename, stylesheet=stylesheet, parser=etree.HTMLParser())
80
    assert "WARNING" in caplog.text
81