Refactor ESPHome camera to avoid creating tasks

homeassistant/components/esphome/camera.py

@@ -2,6 +2,8 @@
 from __future__ import annotations
 
 import asyncio
+from collections.abc import Callable, Coroutine
+from functools import partial
 from typing import Any
 
 from aioesphomeapi import CameraInfo, CameraState
@@ -40,48 +42,56 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
         """Initialize."""
         Camera.__init__(self)
         EsphomeEntity.__init__(self, *args, **kwargs)
-        self._image_cond = asyncio.Condition()
+        self._loop = asyncio.get_running_loop()
+        self._image_futures: list[asyncio.Future[bool | None]] = []
+
+    @callback
+    def _set_futures(self, result: bool) -> None:
+        """Set futures to done."""
+        for future in self._image_futures:
+            if not future.done():
+                future.set_result(result)
+        self._image_futures.clear()
+
+    @callback
+    def _on_device_update(self) -> None:
+        """Handle device going available or unavailable."""
+        super()._on_device_update()
+        if not self.available:
+            self._set_futures(False)
 
     @callback
     def _on_state_update(self) -> None:
         """Notify listeners of new image when update arrives."""
         super()._on_state_update()
-        self.hass.async_create_task(self._on_state_update_coro())
-
-    async def _on_state_update_coro(self) -> None:
-        async with self._image_cond:
-            self._image_cond.notify_all()
+        self._set_futures(True)
 
     async def async_camera_image(
         self, width: int | None = None, height: int | None = None
     ) -> bytes | None:
         """Return single camera image bytes."""
-        if not self.available:
-            return None
-        await self._client.request_single_image()
-        async with self._image_cond:
-            await self._image_cond.wait()
-            if not self.available:
-                # Availability can change while waiting for 'self._image.cond'
-                return None  # type: ignore[unreachable]
-            return self._state.data[:]
+        return await self._async_request_image(self._client.request_single_image)
 
-    async def _async_camera_stream_image(self) -> bytes | None:
-        """Return a single camera image in a stream."""
+    async def _async_request_image(
+        self, request_method: Callable[[], Coroutine[Any, Any, None]]
+    ) -> bytes | None:
+        """Wait for an image to be available and return it."""
         if not self.available:
             return None
-        await self._client.request_image_stream()
-        async with self._image_cond:
-            await self._image_cond.wait()
-            if not self.available:
-                # Availability can change while waiting for 'self._image.cond'
-                return None  # type: ignore[unreachable]
-            return self._state.data[:]
+        image_future = self._loop.create_future()
+        self._image_futures.append(image_future)
+        await request_method()
+        if not await image_future:
+            return None
+        return self._state.data
 
     async def handle_async_mjpeg_stream(
         self, request: web.Request
     ) -> web.StreamResponse:
         """Serve an HTTP MJPEG stream from the camera."""
+        stream_request = partial(
+            self._async_request_image, self._client.request_image_stream
+        )
         return await camera.async_get_still_stream(
-            request, self._async_camera_stream_image, camera.DEFAULT_CONTENT_TYPE, 0.0
+            request, stream_request, camera.DEFAULT_CONTENT_TYPE, 0.0
         )

tests/components/esphome/test_camera.py

@@ -149,9 +149,6 @@ async def test_camera_single_image_unavailable_during_request(
 
     async def _mock_camera_image():
         await mock_device.mock_disconnect(False)
-        # Currently there is a bug where the camera will block
-        # forever if we don't send a response
-        mock_device.set_state(CameraState(key=1, data=SMALLEST_VALID_JPEG_BYTES))
 
     mock_client.request_single_image = _mock_camera_image