From edee57d4e2fe64d904018640b3b5ece084adbae9 Mon Sep 17 00:00:00 2001 From: Erik Montnemery Date: Wed, 2 Sep 2026 17:05:04 +0200 Subject: [PATCH] Improve error handling in file_upload (#181078) --- .../components/file_upload/__init__.py | 50 +++++--- tests/components/file_upload/test_init.py | 116 +++++++++++++++++- 2 files changed, 144 insertions(+), 22 deletions(-) diff --git a/homeassistant/components/file_upload/__init__.py b/homeassistant/components/file_upload/__init__.py index 067e87f559f7..b535b080767d 100644 --- a/homeassistant/components/file_upload/__init__.py +++ b/homeassistant/components/file_upload/__init__.py @@ -176,27 +176,37 @@ class FileUploadView(HomeAssistantView): fut: asyncio.Future[None] | None = None try: - fut = hass.async_add_executor_job(_sync_queue_consumer) - megabytes_sending = 0 - while chunk := await file_field_reader.read_chunk(ONE_MEGABYTE): - megabytes_sending += 1 - if megabytes_sending % 5 != 0: - queue.put_nowait((chunk, None)) - continue + try: + fut = hass.async_add_executor_job(_sync_queue_consumer) + chunks_sent = 0 + while chunk := await file_field_reader.read_chunk(ONE_MEGABYTE): + chunks_sent += 1 + if chunks_sent % 5 != 0: + queue.put_nowait((chunk, None)) + continue - chunk_future = hass.loop.create_future() - queue.put_nowait((chunk, chunk_future)) - await asyncio.wait( - (fut, chunk_future), return_when=asyncio.FIRST_COMPLETED - ) - if fut.done(): - # The executor job failed - break - - queue.put_nowait(None) # terminate queue consumer - finally: - if fut is not None: - await fut + chunk_future = hass.loop.create_future() + queue.put_nowait((chunk, chunk_future)) + await asyncio.wait( + (fut, chunk_future), return_when=asyncio.FIRST_COMPLETED + ) + if fut.done(): + # The executor job failed + break + finally: + # Always terminate the queue consumer, also if the stream raised or + # the task was cancelled. + queue.put_nowait(None) + if fut is not None: + await fut + except Exception, asyncio.CancelledError: + # Upload failed: the consumer has finished and closed the file (inner + # finally above), so removing the directory now cannot race the writer. + # ignore_errors covers a failure that happened before the dir was created. + await hass.async_add_executor_job( + lambda: shutil.rmtree(file_dir, ignore_errors=True) + ) + raise file_upload_data.files[file_id] = filename diff --git a/tests/components/file_upload/test_init.py b/tests/components/file_upload/test_init.py index 49daf1b3b63f..6f940e63fc07 100644 --- a/tests/components/file_upload/test_init.py +++ b/tests/components/file_upload/test_init.py @@ -1,15 +1,19 @@ """Test the File Upload integration.""" +import asyncio from contextlib import contextmanager +from io import StringIO from pathlib import Path from random import getrandbits from typing import Any -from unittest.mock import patch +from unittest.mock import AsyncMock, patch +from aiohttp import BodyPartReader import pytest from homeassistant.components import file_upload -from homeassistant.components.file_upload import DOMAIN +from homeassistant.components.file_upload import DOMAIN, FileUploadView +from homeassistant.components.http import KEY_HASS from homeassistant.core import HomeAssistant from homeassistant.setup import async_setup_component @@ -170,3 +174,111 @@ async def test_upload_large_file_fails( response = await res.content.read() assert b"Boom" in response + + +async def test_upload_stream_error_releases_lock( + hass: HomeAssistant, + hass_client: ClientSessionGenerator, + large_file_io: StringIO, +) -> None: + """Test a mid-upload stream error propagates and releases the upload lock. + + Models a client disconnect: aiohttp raises from read_chunk when the + connection is lost mid-transfer. If the queue consumer's terminating sentinel + is skipped on the error path, awaiting the consumer deadlocks while holding + the upload lock, wedging every later upload; this guards that regression. + """ + assert await async_setup_component(hass, DOMAIN, {}) + client = await hass_client() + + with ( + patch( + # Patch temp dir name to avoid tests fail running in parallel + "homeassistant.components.file_upload.TEMP_DIR_NAME", + file_upload.TEMP_DIR_NAME + f"-{getrandbits(10):03x}", + ), + patch.object( + BodyPartReader, + "read_chunk", + AsyncMock( + side_effect=[b"partial", ConnectionResetError("Connection lost")] + ), + ), + ): + # Bound the request so a reintroduced deadlock fails fast instead of hanging + async with asyncio.timeout(10): + res = await client.post("/api/file_upload", data={"file": large_file_io}) + + assert res.status == 500 + + # The failed upload must not leave a partially written file orphaned on disk + file_upload_data = hass.data[file_upload.DOMAIN] + assert list(file_upload_data.temp_dir.iterdir()) == [] + + # The upload lock must have been released: a subsequent normal upload succeeds + large_file_io.seek(0) + with patch( + "homeassistant.components.file_upload.TEMP_DIR_NAME", + file_upload.TEMP_DIR_NAME + f"-{getrandbits(10):03x}", + ): + async with asyncio.timeout(10): + res = await client.post("/api/file_upload", data={"file": large_file_io}) + + assert res.status == 200 + + +async def test_upload_cancelled_releases_consumer(hass: HomeAssistant) -> None: + """Test cancelling an upload mid-transfer does not deadlock the consumer. + + Driven at the view level because the test HTTP client cannot produce a true + task cancellation mid-request. Without delivering the queue sentinel on the + cancellation path, awaiting the consumer future would hang forever. + """ + assert await async_setup_component(hass, DOMAIN, {}) + view = FileUploadView() + + first_chunk_sent = asyncio.Event() + blocked = asyncio.Event() # never set, so the stream blocks until cancelled + + class _BlockingPart: + """Fake BodyPartReader that blocks after yielding one chunk.""" + + name = "file" + filename = "blocking.bin" + + async def read_chunk(self, size: int) -> bytes: + if first_chunk_sent.is_set(): + await blocked.wait() + return b"" + first_chunk_sent.set() + return b"chunk" + + part = _BlockingPart() + + class _Reader: + async def next(self) -> _BlockingPart: + return part + + class _Request: + app = {KEY_HASS: hass} + + async def multipart(self) -> _Reader: + return _Reader() + + with ( + patch( + "homeassistant.components.file_upload.TEMP_DIR_NAME", + file_upload.TEMP_DIR_NAME + f"-{getrandbits(10):03x}", + ), + patch("homeassistant.components.file_upload.BodyPartReader", _BlockingPart), + ): + task = asyncio.create_task(view._upload_file(_Request())) + await first_chunk_sent.wait() + task.cancel() + # asyncio.wait does not cancel on timeout, so a deadlocked task stays + # pending and the assertion fails fast instead of the run hanging. + _done, pending = await asyncio.wait({task}, timeout=10) + + assert not pending + with pytest.raises(asyncio.CancelledError): + task.result()