Files
drhider/tests/test_extractor.py
T

138 lines
5.7 KiB
Python

"""
Тесты extractor.py — извлечение текста, ZIP-распаковка.
Запуск: python3 tests/test_extractor.py
"""
import io
import zipfile
import sys
import os
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from drhider.extractor import extract_text, expand_zips, docx_to_markdown, pdf_to_markdown
passed = 0
failed = 0
def check(name, condition):
global passed, failed
if condition:
passed += 1
print(f" ✅ {name}")
else:
failed += 1
print(f" ❌ {name}")
# ═══════════════════════════════════════════════════════════════
# extract_text
# ═══════════════════════════════════════════════════════════════
print("\n=== extract_text ===")
# 1. .txt
text = extract_text("test.txt", "Просто текст".encode("utf-8"))
check(".txt — текст извлечён", text == "Просто текст")
# 2. .txt с кириллицей
text = extract_text("документ.txt", "Привет мир".encode("utf-8"))
check(".txt — кириллица", "Привет" in text)
# 3. .doc (бинарный — пропускается)
text = extract_text("старый.doc", b"\xD0\xCF\x11\xE0\x00" * 10)
check(".doc — заглушка", "not supported" in text.lower() or "binary" in text.lower())
# 4. Неизвестное расширение → UTF-8
text = extract_text("file.xyz", "hello".encode("utf-8"))
check(".xyz — как текст", text == "hello")
# ═══════════════════════════════════════════════════════════════
# expand_zips
# ═══════════════════════════════════════════════════════════════
print("\n=== expand_zips ===")
# 1. Не-ZIP файлы проходят без изменений
files = [("test.txt", b"content", "text/plain")]
result = expand_zips(files)
check("не-ZIP — без изменений", len(result) == 1 and result[0][0] == "test.txt")
# 2. ZIP распаковывается
buf = io.BytesIO()
with zipfile.ZipFile(buf, 'w') as zf:
zf.writestr("inner.txt", b"hello")
result = expand_zips([("archive.zip", buf.getvalue(), "application/zip")])
check("ZIP — распакован", len(result) == 1 and result[0][0] == "inner.txt")
check("ZIP — контент верный", result[0][1] == b"hello")
# 3. ZIP с несколькими файлами
buf = io.BytesIO()
with zipfile.ZipFile(buf, 'w') as zf:
zf.writestr("a.txt", b"aaa")
zf.writestr("b.txt", b"bbb")
result = expand_zips([("multi.zip", buf.getvalue(), "")])
check("ZIP — два файла", len(result) == 2)
# 4. Path traversal защита
buf = io.BytesIO()
with zipfile.ZipFile(buf, 'w') as zf:
zf.writestr("../../etc/passwd", b"bad")
result = expand_zips([("evil.zip", buf.getvalue(), "")])
check("path traversal — заблокирован", all("/" not in r[0] and ".." not in r[0] for r in result))
# 5. Директории пропускаются
buf = io.BytesIO()
with zipfile.ZipFile(buf, 'w') as zf:
zf.writestr("dir/", b"")
zf.writestr("dir/file.txt", b"ok")
result = expand_zips([("dirs.zip", buf.getvalue(), "")])
check("директории — пропущены", len(result) == 1)
# ═══════════════════════════════════════════════════════════════
# docx_to_markdown
# ═══════════════════════════════════════════════════════════════
print("\n=== docx_to_markdown ===")
try:
from docx import Document
# Создаём тестовый DOCX
doc = Document()
doc.add_heading("Заголовок", level=1)
p = doc.add_paragraph("Обычный текст")
p.add_run(" жирный").bold = True
doc.add_paragraph("")
# Таблица
table = doc.add_table(rows=2, cols=2)
table.rows[0].cells[0].text = "A"
table.rows[0].cells[1].text = "B"
table.rows[1].cells[0].text = "1"
table.rows[1].cells[1].text = "2"
buf = io.BytesIO()
doc.save(buf)
md = docx_to_markdown(buf.getvalue())
check("DOCX — заголовок #", "# Заголовок" in md)
check("DOCX — жирный **", "жирный**" in md and "**" in md)
check("DOCX — таблица", "| A | B |" in md)
check("DOCX — разделитель таблицы", "---|---|" in md)
check("DOCX — данные таблицы", "| 1 | 2 |" in md)
except ImportError:
print(" ⚠️ python-docx не установлен — тесты DOCX пропущены")
# ═══════════════════════════════════════════════════════════════
# Итог
# ═══════════════════════════════════════════════════════════════
print(f"\n{'='*50}")
print(f"Результат: {passed} OK, {failed} FAIL из {passed + failed}")
if failed > 0:
print("❌ ТЕСТЫ ПРОВАЛЕНЫ")
sys.exit(1)
else:
print("✅ ВСЕ ТЕСТЫ ПРОЙДЕНЫ")