diff --git a/main.py b/main.py index 2cbf404..be2ec27 100644 --- a/main.py +++ b/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 diff --git a/tests/smoke/test_batch_progress.py b/tests/smoke/test_batch_progress.py new file mode 100644 index 0000000..abf2cb3 --- /dev/null +++ b/tests/smoke/test_batch_progress.py @@ -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)] diff --git a/tests/unit/test_main_attachments.py b/tests/unit/test_main_attachments.py new file mode 100644 index 0000000..a05c37e --- /dev/null +++ b/tests/unit/test_main_attachments.py @@ -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""), + _UploadWithoutGetValue("a.xml", 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"" + assert paths[1].read_bytes() == 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"