fix: raise ServiceUnavailable for unregistered read_id in _stream_multiplexer

Author: zhixiangli

packages/google-cloud-storage/google/cloud/storage/asyncio/_stream_multiplexer.py

@@ -15,10 +15,12 @@
 from __future__ import annotations
 
 import asyncio
-import grpc
 import logging
 from typing import Awaitable, Callable, Dict, Optional, Set
 
+import grpc
+from google.api_core import exceptions
+
 from google.cloud import _storage_v2
 from google.cloud.storage.asyncio.async_read_object_stream import (
     _AsyncReadObjectStream,
@@ -139,6 +141,10 @@ async def _recv_loop(self) -> None:
                         queue = self._queues.get(read_id)
                         if queue:
                             queues_to_notify.add(queue)
+                        else:
+                            raise exceptions.ServiceUnavailable(
+                                f"Received data for unregistered read_id: {read_id}"
+                            )
                     await asyncio.gather(
                         *(
                             self._put_with_timeout(queue, response)

packages/google-cloud-storage/tests/unit/asyncio/test_stream_multiplexer.py

@@ -13,10 +13,11 @@
 # limitations under the License.
 
 import asyncio
-import grpc
 from unittest.mock import AsyncMock
 
+import grpc
 import pytest
+from google.api_core import exceptions
 
 from google.cloud import _storage_v2
 from google.cloud.storage.asyncio._stream_multiplexer import (
@@ -325,7 +326,7 @@ async def test_error_uses_put_nowait(self):
         assert err.exception is exc
 
     @pytest.mark.asyncio
-    async def test_unknown_read_id_is_dropped(self):
+    async def test_unknown_read_id_raises_service_unavailable(self):
         # Given a multiplexer and a response with an unknown read ID
         mock_stream = AsyncMock()
         resp = _make_response(read_id=999)
@@ -337,9 +338,11 @@ async def test_unknown_read_id_is_dropped(self):
         # When the receive loop is started
         mux._ensure_recv_loop()
 
-        # Then the response is dropped and only StreamEnd is received
-        end = await asyncio.wait_for(queue.get(), timeout=1)
-        assert isinstance(end, _StreamEnd)
+        # Then a StreamError is broadcast to all queues
+        err = await asyncio.wait_for(queue.get(), timeout=1)
+        assert isinstance(err, _StreamError)
+        assert isinstance(err.exception, exceptions.ServiceUnavailable)
+        assert "Received data for unregistered read_id: 999" in str(err.exception)
 
 
 class TestSend: