Fix AsyncIteratorReader blocking after stream exhaustion (#161731)

homeassistant/util/async_iterator.py

@@ -28,6 +28,7 @@ class AsyncIteratorReader:
     ) -> None:
         """Initialize the wrapper."""
         self._aborted = False
+        self._exhausted = False
         self._loop = loop
         self._stream = stream
         self._buffer: bytes | None = None
@@ -51,6 +52,8 @@ class AsyncIteratorReader:
         """
         result = bytearray()
         while n < 0 or len(result) < n:
+            if self._exhausted:
+                break
             if not self._buffer:
                 self._next_future = asyncio.run_coroutine_threadsafe(
                     self._next(), self._loop
@@ -65,6 +68,7 @@ class AsyncIteratorReader:
                 self._pos = 0
             if not self._buffer:
                 # The stream is exhausted
+                self._exhausted = True
                 break
             chunk = self._buffer[self._pos : self._pos + n]
             result.extend(chunk)

tests/util/test_async_iterator.py

@@ -114,3 +114,20 @@ async def test_async_iterator_writer_abort_late(hass: HomeAssistant) -> None:
 
     with pytest.raises(Abort):
         await fut
+
+
+async def test_async_iterator_reader_exhausted(hass: HomeAssistant) -> None:
+    """Test that read() returns empty bytes after stream exhaustion."""
+
+    async def async_gen() -> AsyncIterator[bytes]:
+        yield b"hello"
+
+    reader = AsyncIteratorReader(hass.loop, async_gen())
+
+    def _read_then_read_again() -> bytes:
+        data = _read_all(reader)
+        # Second read after exhaustion should return b"" immediately
+        assert reader.read(500) == b""
+        return data
+
+    assert await hass.async_add_executor_job(_read_then_read_again) == b"hello"