AutoGPT

test_text_file_parsers.py
170 строк · 4.8 Кб
Перенос по словам
1
import json
2
import logging
3
import os.path
4
import tempfile
5
from pathlib import Path
6
from xml.etree import ElementTree
7

8
import docx
9
import pytest
10
import yaml
11
from bs4 import BeautifulSoup
12

13
from autogpt.commands.file_operations_utils import (
14
    decode_textual_file,
15
    is_file_binary_fn,
16
)
17

18
logger = logging.getLogger(__name__)
19

20
plain_text_str = "Hello, world!"
21

22

23
def mock_text_file():
24
    with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".txt") as f:
25
        f.write(plain_text_str)
26
    return f.name
27

28

29
def mock_csv_file():
30
    with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".csv") as f:
31
        f.write(plain_text_str)
32
    return f.name
33

34

35
def mock_pdf_file():
36
    with tempfile.NamedTemporaryFile(mode="wb", delete=False, suffix=".pdf") as f:
37
        # Create a new PDF and add a page with the text plain_text_str
38
        # Write the PDF header
39
        f.write(b"%PDF-1.7\n")
40
        # Write the document catalog
41
        f.write(b"1 0 obj\n")
42
        f.write(b"<< /Type /Catalog /Pages 2 0 R >>\n")
43
        f.write(b"endobj\n")
44
        # Write the page object
45
        f.write(b"2 0 obj\n")
46
        f.write(
47
            b"<< /Type /Page /Parent 1 0 R /Resources << /Font << /F1 3 0 R >> >> "
48
            b"/MediaBox [0 0 612 792] /Contents 4 0 R >>\n"
49
        )
50
        f.write(b"endobj\n")
51
        # Write the font object
52
        f.write(b"3 0 obj\n")
53
        f.write(
54
            b"<< /Type /Font /Subtype /Type1 /Name /F1 /BaseFont /Helvetica-Bold >>\n"
55
        )
56
        f.write(b"endobj\n")
57
        # Write the page contents object
58
        f.write(b"4 0 obj\n")
59
        f.write(b"<< /Length 25 >>\n")
60
        f.write(b"stream\n")
61
        f.write(b"BT\n/F1 12 Tf\n72 720 Td\n(Hello, world!) Tj\nET\n")
62
        f.write(b"endstream\n")
63
        f.write(b"endobj\n")
64
        # Write the cross-reference table
65
        f.write(b"xref\n")
66
        f.write(b"0 5\n")
67
        f.write(b"0000000000 65535 f \n")
68
        f.write(b"0000000017 00000 n \n")
69
        f.write(b"0000000073 00000 n \n")
70
        f.write(b"0000000123 00000 n \n")
71
        f.write(b"0000000271 00000 n \n")
72
        f.write(b"trailer\n")
73
        f.write(b"<< /Size 5 /Root 1 0 R >>\n")
74
        f.write(b"startxref\n")
75
        f.write(b"380\n")
76
        f.write(b"%%EOF\n")
77
        f.write(b"\x00")
78
    return f.name
79

80

81
def mock_docx_file():
82
    with tempfile.NamedTemporaryFile(mode="wb", delete=False, suffix=".docx") as f:
83
        document = docx.Document()
84
        document.add_paragraph(plain_text_str)
85
        document.save(f.name)
86
    return f.name
87

88

89
def mock_json_file():
90
    with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".json") as f:
91
        json.dump({"text": plain_text_str}, f)
92
    return f.name
93

94

95
def mock_xml_file():
96
    root = ElementTree.Element("text")
97
    root.text = plain_text_str
98
    tree = ElementTree.ElementTree(root)
99
    with tempfile.NamedTemporaryFile(mode="wb", delete=False, suffix=".xml") as f:
100
        tree.write(f)
101
    return f.name
102

103

104
def mock_yaml_file():
105
    with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".yaml") as f:
106
        yaml.dump({"text": plain_text_str}, f)
107
    return f.name
108

109

110
def mock_html_file():
111
    html = BeautifulSoup(
112
        "<html>"
113
        "<head><title>This is a test</title></head>"
114
        f"<body><p>{plain_text_str}</p></body>"
115
        "</html>",
116
        "html.parser",
117
    )
118
    with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".html") as f:
119
        f.write(str(html))
120
    return f.name
121

122

123
def mock_md_file():
124
    with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".md") as f:
125
        f.write(f"# {plain_text_str}!\n")
126
    return f.name
127

128

129
def mock_latex_file():
130
    with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".tex") as f:
131
        latex_str = (
132
            r"\documentclass{article}"
133
            r"\begin{document}"
134
            f"{plain_text_str}"
135
            r"\end{document}"
136
        )
137
        f.write(latex_str)
138
    return f.name
139

140

141
respective_file_creation_functions = {
142
    ".txt": mock_text_file,
143
    ".csv": mock_csv_file,
144
    ".pdf": mock_pdf_file,
145
    ".docx": mock_docx_file,
146
    ".json": mock_json_file,
147
    ".xml": mock_xml_file,
148
    ".yaml": mock_yaml_file,
149
    ".html": mock_html_file,
150
    ".md": mock_md_file,
151
    ".tex": mock_latex_file,
152
}
153
binary_files_extensions = [".pdf", ".docx"]
154

155

156
@pytest.mark.parametrize(
157
    "file_extension, c_file_creator",
158
    respective_file_creation_functions.items(),
159
)
160
def test_parsers(file_extension, c_file_creator):
161
    created_file_path = Path(c_file_creator())
162
    with open(created_file_path, "rb") as file:
163
        loaded_text = decode_textual_file(file, os.path.splitext(file.name)[1], logger)
164

165
        assert plain_text_str in loaded_text
166

167
        should_be_binary = file_extension in binary_files_extensions
168
        assert should_be_binary == is_file_binary_fn(file)
169

170
    created_file_path.unlink()  # cleanup
171
AutoGPT

Использование cookies
