138 lines
4.8 KiB
Python
138 lines
4.8 KiB
Python
"""Сохранение вложений Streamlit на диск не должно удваивать их в памяти.
|
||
|
||
UploadedFile — это BytesIO: Streamlit уже держит всё содержимое вложения в
|
||
памяти. getvalue() делает ещё одну полную копию каждого файла перед записью
|
||
на диск; при 230 вложениях в контейнере на 1 ГБ это лишняя пиковая нагрузка.
|
||
Копировать нужно потоково, через файловый интерфейс, не трогая getvalue().
|
||
|
||
Исключение вложений из обработки (без удаления из file_uploader) и сброс
|
||
загрузчика — по file_id, потому что имена вложений могут совпадать.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import io
|
||
from pathlib import Path
|
||
|
||
from main import (
|
||
EXCLUDED_FILE_IDS_KEY,
|
||
UPLOADER_GENERATION_KEY,
|
||
_active_uploads,
|
||
_clear_all_attachments,
|
||
_exclude_checkbox_key,
|
||
_exclude_selected,
|
||
_init_attachment_state,
|
||
_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"
|
||
|
||
|
||
class _Upload:
|
||
"""Минимальный UploadedFile: имя и file_id, по которому его отличаем."""
|
||
|
||
def __init__(self, name: str, file_id: str, size: int = 10) -> None:
|
||
self.name = name
|
||
self.file_id = file_id
|
||
self.size = size
|
||
|
||
|
||
def test_excluded_files_are_matched_by_file_id_not_by_name() -> None:
|
||
first = _Upload("Файлы.zip", "id-1")
|
||
second = _Upload("Файлы.zip", "id-2")
|
||
other = _Upload("a.xml", "id-3")
|
||
|
||
active = _active_uploads([first, second, other], {"id-2"})
|
||
|
||
assert active == [first, other]
|
||
|
||
|
||
def test_exclude_selected_moves_checked_files_to_excluded() -> None:
|
||
state: dict = {}
|
||
_init_attachment_state(state)
|
||
first = _Upload("Файлы.zip", "id-1")
|
||
second = _Upload("Файлы.zip", "id-2")
|
||
state[_exclude_checkbox_key(0, "id-2")] = True
|
||
state[_exclude_checkbox_key(0, "id-1")] = False
|
||
|
||
excluded = _exclude_selected(state, [first, second])
|
||
|
||
assert excluded == 1
|
||
assert state[EXCLUDED_FILE_IDS_KEY] == {"id-2"}
|
||
assert _active_uploads([first, second], state[EXCLUDED_FILE_IDS_KEY]) == [first]
|
||
|
||
|
||
def test_exclude_selected_keeps_previous_exclusions() -> None:
|
||
state: dict = {EXCLUDED_FILE_IDS_KEY: {"id-1"}, UPLOADER_GENERATION_KEY: 2}
|
||
state[_exclude_checkbox_key(2, "id-3")] = True
|
||
# Отметка из прошлого поколения загрузчика не должна ничего исключать.
|
||
state[_exclude_checkbox_key(1, "id-4")] = True
|
||
|
||
excluded = _exclude_selected(
|
||
state, [_Upload("c.xml", "id-3"), _Upload("d.xml", "id-4")]
|
||
)
|
||
|
||
assert excluded == 1
|
||
assert state[EXCLUDED_FILE_IDS_KEY] == {"id-1", "id-3"}
|
||
|
||
|
||
def test_exclude_selected_without_checked_files_changes_nothing() -> None:
|
||
state: dict = {}
|
||
_init_attachment_state(state)
|
||
|
||
assert _exclude_selected(state, [_Upload("a.xml", "id-1")]) == 0
|
||
assert state[EXCLUDED_FILE_IDS_KEY] == set()
|
||
|
||
|
||
def test_clear_all_resets_uploader_key_and_exclusions() -> None:
|
||
state: dict = {}
|
||
_init_attachment_state(state)
|
||
state[EXCLUDED_FILE_IDS_KEY] = {"id-1", "id-2"}
|
||
|
||
_clear_all_attachments(state)
|
||
|
||
assert state[UPLOADER_GENERATION_KEY] == 1
|
||
assert state[EXCLUDED_FILE_IDS_KEY] == set()
|
||
|
||
|
||
def test_init_attachment_state_keeps_existing_values() -> None:
|
||
state: dict = {UPLOADER_GENERATION_KEY: 5, EXCLUDED_FILE_IDS_KEY: {"id-1"}}
|
||
|
||
_init_attachment_state(state)
|
||
|
||
assert state == {UPLOADER_GENERATION_KEY: 5, EXCLUDED_FILE_IDS_KEY: {"id-1"}}
|