Ускорить формирование отчёта: обратный индекс маппинга, потоковые вложения, тесты прогресса #2
8
main.py
8
main.py
@ -1,4 +1,5 @@
|
||||
import asyncio
|
||||
import shutil
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
|
||||
@ -122,7 +123,12 @@ def _save_attachments(uploaded_files, target_dir: Path) -> list[Path]:
|
||||
holder = target_dir / str(index)
|
||||
holder.mkdir(parents=True, exist_ok=True)
|
||||
path = holder / uploaded.name
|
||||
path.write_bytes(uploaded.getvalue())
|
||||
# Потоковое копирование: UploadedFile уже целиком лежит в памяти
|
||||
# Streamlit, а getvalue() делал бы вторую полную копию каждого
|
||||
# вложения перед записью на диск.
|
||||
uploaded.seek(0)
|
||||
with path.open("wb") as target:
|
||||
shutil.copyfileobj(uploaded, target)
|
||||
paths.append(path)
|
||||
return paths
|
||||
|
||||
|
||||
59
tests/smoke/test_batch_progress.py
Normal file
59
tests/smoke/test_batch_progress.py
Normal file
@ -0,0 +1,59 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import zipfile
|
||||
from pathlib import Path
|
||||
|
||||
from app.pipeline import service
|
||||
from app.pipeline.service import run_batch
|
||||
from tests.smoke.test_batch_smoke import _patch_opencloud, _valid_xml
|
||||
|
||||
|
||||
def test_progress_callback_is_called_once_per_xml_file(
|
||||
tmp_path: Path, monkeypatch
|
||||
) -> None:
|
||||
_patch_opencloud(monkeypatch)
|
||||
zip_path = tmp_path / "batch.zip"
|
||||
names = [f"SKO115FZ_01_123456789_20260616_X0000{i}.xml" for i in range(1, 4)]
|
||||
with zipfile.ZipFile(zip_path, "w") as archive:
|
||||
for name in names:
|
||||
archive.writestr(name, _valid_xml())
|
||||
loose = tmp_path / "SKO115FZ_01_123456789_20260616_X00009.xml"
|
||||
loose.write_text(_valid_xml(), encoding="utf-8")
|
||||
|
||||
calls: list[tuple[int, int, str]] = []
|
||||
result = run_batch(
|
||||
input_paths=[zip_path, loose],
|
||||
launcher="tester",
|
||||
opencloud_username="test-user",
|
||||
opencloud_access_token="token-123",
|
||||
progress_callback=lambda done, total, name: calls.append((done, total, name)),
|
||||
)
|
||||
|
||||
assert result.status == "success"
|
||||
assert [done for done, _, _ in calls] == [1, 2, 3, 4]
|
||||
assert {total for _, total, _ in calls} == {4}
|
||||
assert sorted(name for _, _, name in calls) == sorted([*names, loose.name])
|
||||
|
||||
|
||||
def test_progress_callback_is_called_for_unreadable_file_too(
|
||||
tmp_path: Path, monkeypatch
|
||||
) -> None:
|
||||
_patch_opencloud(monkeypatch)
|
||||
loose = tmp_path / "SKO115FZ_01_123456789_20260616_X00001.xml"
|
||||
loose.write_text(_valid_xml(), encoding="utf-8")
|
||||
|
||||
def _boom(self: Path) -> bytes:
|
||||
raise OSError("disk gone")
|
||||
|
||||
monkeypatch.setattr(service.Path, "read_bytes", _boom)
|
||||
calls: list[tuple[int, int, str]] = []
|
||||
result = run_batch(
|
||||
input_paths=[loose],
|
||||
launcher="tester",
|
||||
opencloud_username="test-user",
|
||||
opencloud_access_token="token-123",
|
||||
progress_callback=lambda done, total, name: calls.append((done, total, name)),
|
||||
)
|
||||
|
||||
assert result.status == "error"
|
||||
assert calls == [(1, 1, loose.name)]
|
||||
50
tests/unit/test_main_attachments.py
Normal file
50
tests/unit/test_main_attachments.py
Normal file
@ -0,0 +1,50 @@
|
||||
"""Сохранение вложений Streamlit на диск не должно удваивать их в памяти.
|
||||
|
||||
UploadedFile — это BytesIO: Streamlit уже держит всё содержимое вложения в
|
||||
памяти. getvalue() делает ещё одну полную копию каждого файла перед записью
|
||||
на диск; при 230 вложениях в контейнере на 1 ГБ это лишняя пиковая нагрузка.
|
||||
Копировать нужно потоково, через файловый интерфейс, не трогая getvalue().
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import io
|
||||
from pathlib import Path
|
||||
|
||||
from main import _save_attachments
|
||||
|
||||
|
||||
class _UploadWithoutGetValue(io.BytesIO):
|
||||
"""Вложение, у которого есть только файловый интерфейс и имя."""
|
||||
|
||||
def __init__(self, name: str, payload: bytes) -> None:
|
||||
super().__init__(payload)
|
||||
self.name = name
|
||||
|
||||
def getvalue(self) -> bytes: # type: ignore[override]
|
||||
raise AssertionError("getvalue() копирует всё вложение в память")
|
||||
|
||||
|
||||
def test_save_attachments_streams_without_getvalue(tmp_path: Path) -> None:
|
||||
uploads = [
|
||||
_UploadWithoutGetValue("a.xml", b"<a/>"),
|
||||
_UploadWithoutGetValue("a.xml", b"<b/>" * 1000),
|
||||
]
|
||||
|
||||
paths = _save_attachments(uploads, tmp_path)
|
||||
|
||||
assert [path.name for path in paths] == ["a.xml", "a.xml"]
|
||||
assert paths[0] != paths[1]
|
||||
assert paths[0].read_bytes() == b"<a/>"
|
||||
assert paths[1].read_bytes() == b"<b/>" * 1000
|
||||
|
||||
|
||||
def test_save_attachments_copies_from_start_even_if_upload_was_read(
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
upload = _UploadWithoutGetValue("f.xml", b"0123456789")
|
||||
upload.read(4)
|
||||
|
||||
(path,) = _save_attachments([upload], tmp_path)
|
||||
|
||||
assert path.read_bytes() == b"0123456789"
|
||||
Loading…
x
Reference in New Issue
Block a user