diff --git a/src/labapi/entry/entries/attachment.py b/src/labapi/entry/entries/attachment.py index 05c40b1..df1425f 100644 --- a/src/labapi/entry/entries/attachment.py +++ b/src/labapi/entry/entries/attachment.py @@ -6,6 +6,7 @@ from __future__ import annotations +import shutil from email.message import Message from io import BytesIO from tempfile import TemporaryFile @@ -89,7 +90,7 @@ def get_attachment(self, use_tempfile: bool = False) -> Attachment: # Return an independent copy so each caller gets isolated read/seek/close state # while still sharing a single downloaded backing attachment in the cache. self._filedata.seek(0) - output.write(self._filedata.read()) + shutil.copyfileobj(self._filedata, output) output.seek(0) return Attachment( diff --git a/tests/entry/entries/test_attachment.py b/tests/entry/entries/test_attachment.py index 51b0da0..f89f292 100644 --- a/tests/entry/entries/test_attachment.py +++ b/tests/entry/entries/test_attachment.py @@ -169,3 +169,23 @@ def test_attachment_entry_get_attachment_caching(self, client, user: User): attachment3 = entry.get_attachment() assert client.stream_api_get.call_count == 1 assert attachment3.read() == b"Content" + + def test_get_attachment_tempfile_copies_in_chunks(self, client, user: User): + """get_attachment(use_tempfile=True) must not read the full payload at once. + + Regression for #215: the old code called self._filedata.read() with no + size, materializing the entire cached payload as one bytes object. + """ + entry = AttachmentEntry("eid_att", "Caption", user) + + mock_response = Mock() + mock_response.headers = { + "Content-Type": "text/plain", + "Content-Disposition": 'attachment; filename="test.txt"', + } + mock_response.iter_content.return_value = [b"Chunked content"] + client.stream_api_get = Mock(return_value=StreamingResponse(mock_response)) + + attachment = entry.get_attachment(use_tempfile=True) + assert attachment.read() == b"Chunked content" + attachment.close()