e45502b2ea
- ClosingStreamingResponse закрывает body_iterator в finally (под shield), ChainedStream закрывает исходный поток и до начала чтения; заголовок X-Content-Type-Options: nosniff; - ASCII-имя в Content-Disposition строится из file_type, а не из расширения исходного имени; - ObjectStorage.stream типизирован как AsyncGenerator, убраны два type: ignore. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_012vJkmBGKKqMoEMzvsXG3y2
222 lines
9.9 KiB
Python
222 lines
9.9 KiB
Python
"""Исходный файл потоком (spec documents-api «Исходный файл»; design D-10)."""
|
||
|
||
from collections.abc import AsyncGenerator
|
||
|
||
import pytest
|
||
from httpx import AsyncClient
|
||
from sqlalchemy import select
|
||
from sqlalchemy.ext.asyncio import AsyncSession
|
||
|
||
from src.common.streaming import ClosingStreamingResponse
|
||
from src.core.storage import ObjectNotFoundError
|
||
from src.models.enums import FileStatus, FileType, RecognitionStatus
|
||
from src.models.files import S3Object, StoredFile
|
||
from src.services.primary_files import _chain
|
||
from tests.conftest import AuthClientFactory, FakeStorage
|
||
from tests.factories import make_stored_file
|
||
from tests.test_api_documents import CLIENT_MANAGER, CURATOR_GO, KA_RF, RKA, _rostov_entity
|
||
|
||
pytestmark = pytest.mark.integration
|
||
PDF = b"%PDF-1.4\n" + bytes(range(256)) * 4096 # ~1 МиБ: проверка порционной отдачи
|
||
HTTP_OK = 200
|
||
HTTP_NOT_FOUND = 404
|
||
HTTP_FORBIDDEN = 403
|
||
SECOND_BODY = 3
|
||
|
||
|
||
async def _file(
|
||
session: AsyncSession, storage: FakeStorage, name: str, file_type: FileType, *, entity_name: str = "ООО Файлы"
|
||
) -> StoredFile:
|
||
entity = await _rostov_entity(session, name=entity_name)
|
||
file, _ = await make_stored_file(
|
||
session, entity=entity, status=FileStatus.PROCESSED, job_status=RecognitionStatus.done, original_name=name
|
||
)
|
||
file.file_type = file_type
|
||
key = (await session.execute(select(S3Object.key).where(S3Object.id == file.s3_object_id))).scalar_one()
|
||
storage.objects[key] = PDF
|
||
await session.flush()
|
||
return file
|
||
|
||
|
||
async def test_inline_pdf(app, seeded: AsyncSession, auth_client: AuthClientFactory) -> None: # noqa: ANN001
|
||
storage = FakeStorage()
|
||
app.state.storage = storage
|
||
file = await _file(seeded, storage, "aep_6.pdf", FileType.PDF)
|
||
client: AsyncClient = await auth_client(KA_RF)
|
||
response = await client.get(f"/files/{file.id}/content", params={"disposition": "inline"})
|
||
assert response.status_code == HTTP_OK
|
||
assert response.headers["content-type"] == "application/pdf"
|
||
assert response.headers["content-disposition"] == "inline; filename=\"file.pdf\"; filename*=UTF-8''aep_6.pdf"
|
||
assert response.content == PDF
|
||
assert response.headers["x-content-type-options"] == "nosniff" # T9.5
|
||
|
||
|
||
async def test_cyrillic_name_and_missing(app, seeded: AsyncSession, auth_client: AuthClientFactory) -> None: # noqa: ANN001
|
||
storage = FakeStorage()
|
||
app.state.storage = storage
|
||
file = await _file(seeded, storage, "ОСВ_66 за 2024.xls", FileType.XLS)
|
||
client = await auth_client(KA_RF)
|
||
response = await client.get(f"/files/{file.id}/content")
|
||
assert response.headers["content-type"] == "application/vnd.ms-excel"
|
||
assert response.headers["content-disposition"].startswith(
|
||
"attachment; filename=\"file.xls\"; filename*=UTF-8''%D0%9E"
|
||
)
|
||
storage.objects.clear()
|
||
missing = await client.get(f"/files/{file.id}/content")
|
||
assert missing.status_code == HTTP_NOT_FOUND and missing.json()["detail"]["code"] == "FILE_CONTENT_NOT_FOUND"
|
||
|
||
|
||
async def test_superseded_file_is_served(app, seeded: AsyncSession, auth_client: AuthClientFactory) -> None: # noqa: ANN001
|
||
storage = FakeStorage()
|
||
app.state.storage = storage
|
||
file = await _file(seeded, storage, "old.xlsx", FileType.XLSX)
|
||
file.status = FileStatus.SUPERSEDED
|
||
await seeded.flush()
|
||
client = await auth_client(KA_RF)
|
||
response = await client.get(f"/files/{file.id}/content")
|
||
assert response.status_code == HTTP_OK and response.content == PDF
|
||
assert response.headers["content-type"] == ("application/vnd.openxmlformats-officedocument.spreadsheetml.sheet")
|
||
|
||
|
||
async def test_foreign_unknown_and_permission(app, seeded: AsyncSession, auth_client: AuthClientFactory) -> None: # noqa: ANN001
|
||
storage = FakeStorage()
|
||
app.state.storage = storage
|
||
file = await _file(seeded, storage, "a.pdf", FileType.PDF)
|
||
# RKA — только DOCUMENT_VIEW: доступ есть (тот же филиал); без права — 403 до видимости
|
||
assert (await (await auth_client(RKA)).get(f"/files/{file.id}/content")).status_code == HTTP_OK
|
||
assert (await (await auth_client(CURATOR_GO)).get(f"/files/{file.id}/content")).status_code == HTTP_FORBIDDEN
|
||
foreign = await (await auth_client(CLIENT_MANAGER)).get(f"/files/{file.id}/content") # другой филиал
|
||
assert foreign.status_code == HTTP_NOT_FOUND and foreign.json()["detail"]["code"] == "NOT_FOUND"
|
||
client = await auth_client(KA_RF)
|
||
unknown = await client.get("/files/00000000-0000-0000-0000-000000000000/content")
|
||
assert unknown.status_code == HTTP_NOT_FOUND
|
||
|
||
|
||
async def test_stream_error_before_first_chunk_is_object_not_found() -> None:
|
||
storage = FakeStorage()
|
||
with pytest.raises(ObjectNotFoundError):
|
||
await storage.stream("missing").__anext__()
|
||
|
||
|
||
async def test_stream_releases_on_error() -> None:
|
||
"""Review Focus 4: при обрыве потребителя ответ S3 закрывается и отпускается."""
|
||
from src.core.clients import S3Storage # noqa: PLC0415
|
||
|
||
class Content:
|
||
async def iter_chunked(self, size: int): # noqa: ANN202
|
||
yield b"a" * size
|
||
yield b"b" * size
|
||
|
||
class Response:
|
||
def __init__(self) -> None:
|
||
self.content = Content()
|
||
self.closed = self.released = False
|
||
|
||
def close(self) -> None:
|
||
self.closed = True
|
||
|
||
async def release(self) -> None:
|
||
self.released = True
|
||
|
||
response = Response()
|
||
|
||
class Client:
|
||
async def get_object(self, bucket: str, name: str) -> Response:
|
||
return response
|
||
|
||
storage = S3Storage(Client(), "arm-ka") # type: ignore[arg-type]
|
||
stream = storage.stream("uploads/x", chunk_size=4)
|
||
assert await stream.__anext__() == b"aaaa"
|
||
await stream.aclose() # потребитель оборвал чтение
|
||
assert response.closed and response.released
|
||
|
||
|
||
class _Source:
|
||
"""Источник порций с признаком закрытия (`finally` асинхронного генератора)."""
|
||
|
||
def __init__(self, fail_after: int | None = None) -> None:
|
||
self.closed = False
|
||
self._fail_after = fail_after
|
||
self.iterator = self._run()
|
||
|
||
async def _run(self) -> AsyncGenerator[bytes, None]:
|
||
try:
|
||
for index in range(3):
|
||
if self._fail_after == index:
|
||
raise ConnectionError("обрыв S3")
|
||
yield b"x"
|
||
finally:
|
||
self.closed = True
|
||
|
||
|
||
async def test_chain_closes_source_when_consumer_stops() -> None:
|
||
"""Обрыв клиента после первой порции: `aclose` обёртки закрывает исходный поток."""
|
||
source = _Source()
|
||
first = await anext(source.iterator, b"")
|
||
chain = _chain(first, source.iterator)
|
||
assert await chain.__anext__() == b"x"
|
||
await chain.aclose()
|
||
assert source.closed
|
||
|
||
|
||
async def test_chain_closes_source_on_error_and_yields_all() -> None:
|
||
source = _Source(fail_after=2)
|
||
chain = _chain(await anext(source.iterator, b""), source.iterator)
|
||
with pytest.raises(ConnectionError):
|
||
_ = [chunk async for chunk in chain]
|
||
assert source.closed
|
||
complete = _Source()
|
||
assert b"".join([chunk async for chunk in _chain(await anext(complete.iterator, b""), complete.iterator)]) == b"xxx"
|
||
assert complete.closed
|
||
|
||
|
||
async def test_chain_aclose_closes_source_when_never_iterated() -> None:
|
||
"""T9.1: поток не начали отдавать (первая порция уже прочитана) — `aclose` всё равно закрывает S3-поток."""
|
||
source = _Source()
|
||
chain = _chain(await anext(source.iterator, b""), source.iterator)
|
||
await chain.aclose()
|
||
assert source.closed
|
||
|
||
|
||
async def test_response_closes_stream_when_client_drops_after_first_chunk() -> None:
|
||
"""T9.1: `send` падает на втором теле (обрыв клиента) — `body_iterator` закрыт, S3-поток освобождён."""
|
||
source = _Source()
|
||
chain = _chain(await anext(source.iterator, b""), source.iterator)
|
||
response = ClosingStreamingResponse(chain, media_type="application/pdf")
|
||
sent: list[str] = []
|
||
|
||
async def send(message: dict) -> None: # type: ignore[type-arg]
|
||
sent.append(message["type"])
|
||
if len(sent) == SECOND_BODY: # start, первая порция, вторая порция
|
||
raise OSError("клиент отключился")
|
||
|
||
with pytest.raises(OSError, match="клиент"):
|
||
await response.stream_response(send)
|
||
assert source.closed
|
||
|
||
|
||
async def test_response_closes_stream_when_send_fails_at_start() -> None:
|
||
"""T9.1: тело не начали итерировать (падает уже отправка заголовков) — поток тоже закрыт."""
|
||
source = _Source()
|
||
chain = _chain(await anext(source.iterator, b""), source.iterator)
|
||
response = ClosingStreamingResponse(chain)
|
||
|
||
async def send(message: dict) -> None: # type: ignore[type-arg]
|
||
raise OSError("клиент отключился")
|
||
|
||
with pytest.raises(OSError, match="клиент"):
|
||
await response.stream_response(send)
|
||
assert source.closed
|
||
|
||
|
||
async def test_ascii_extension_from_file_type(app, seeded: AsyncSession, auth_client: AuthClientFactory) -> None: # noqa: ANN001
|
||
"""T9.2: ASCII-имя берётся из `file_type`, а не из расширения исходного имени (нестандартное/кириллическое)."""
|
||
storage = FakeStorage()
|
||
app.state.storage = storage
|
||
file = await _file(seeded, storage, "отчёт.экс", FileType.XLSX)
|
||
client = await auth_client(KA_RF)
|
||
response = await client.get(f"/files/{file.id}/content")
|
||
disposition = response.headers["content-disposition"]
|
||
assert disposition.startswith("attachment; filename=\"file.xlsx\"; filename*=UTF-8''")
|
||
assert disposition.isascii()
|