update response iterator from async auth

Author: ohmayr

google/api_core/rest_streaming_async.py

@@ -17,8 +17,8 @@
 from typing import Union
 
 import proto
-import google.protobuf.message
 import google.auth.aio.transport
+import google.protobuf.message
 from google.api_core._rest_streaming_base import BaseResponseIterator
 
 
@@ -43,58 +43,18 @@ def __init__(
     ):
         self._response = response
         self._chunk_size = 1024
-        self._response_itr = None
+        self._response_itr = self._response.content().__aiter__()
         super(AsyncResponseIterator, self).__init__(
             response_message_cls=response_message_cls
         )
 
-    async def _create_response_iter(self):
-
-        if not isinstance(self._response, google.auth.aio.transport.Response):
-            raise ValueError(
-                "Response must be of type google.auth.aio.transport.Response"
-            )
-
-        # TODO (ohmayr): Ideally the response from auth should expose an attribute
-        # to read streaming response iterator directly i.e.
-        #
-        # self -> google.auth.aio.transport.aiohttp.Response:
-        # def stream_response_itr(self, chunk_size):
-        #   return self._response.content.iter_chunked(chunk_size)
-        #
-        # self -> google.auth.aio.transport.httpx.Response:
-        # def stream_response_itr(self, chunk_size):
-        #   return self._response.aiter_raw(chunk_size)
-        #
-        # this way we can just call the property directly to get the appropriate
-        # response iterator without having to deal with the underlying API differences
-        # or alternatively, having to check the type of inherited response types here
-        # i.e we could do: self._response_itr = self._response.stream_response_itr(self._chunk_size)
-
-        content = self._response.content
-        if hasattr(content, "iter_chunked"):
-            return content.iter_chunked(self._chunk_size)
-        else:
-            # TODO (ohmayr): since iter_chunked is only available in an instance of
-            # google.auth.aio.transport.aiohttp.Response, we indirectly depend on
-            # on the inherited class.
-            raise ValueError(
-                f"Unsupported Response type: {type(self._response)}. Expected google.auth.aio.transport.aiohttp.Response."
-            )
-
     async def cancel(self):
         """Cancel existing streaming operation."""
         await self._response.close()
 
     async def __anext__(self):
         while not self._ready_objs:
-            try:
-                if not self._response_itr:
-                    self._response_itr = await self._create_response_iter()
-                    # TODO (ohmayr): cleanup
-                    # content = await self._response.content
-                    # self._response_itr = content.iter_chunked(self._chunk_size)
-
+            try:    
                 chunk = await self._response_itr.__anext__()
                 chunk = chunk.decode("utf-8")
                 self._process_chunk(chunk)
@@ -110,5 +70,5 @@ def __aiter__(self):
         return self
 
     async def __aexit__(self, exc_type, exc, tb):
-        """Cancel existing streaming operation."""
+        """Cancel existing async streaming operation."""
         await self._response.close()

tests/_helpers.py renamed to tests/conftest.py

@@ -27,8 +27,7 @@
 from google.api import http_pb2
 from google.api import httpbody_pb2
 
-# TODO (ohmayr): confirm if relative path is not an issue.
-from .._helpers import Composer, Song, EchoResponse, parse_responses
+from ..conftest import Composer, Song, EchoResponse, parse_responses
 
 
 __protobuf__ = proto.module(package=__name__)

tests/unit/test_rest_streaming.py

@@ -17,18 +17,17 @@
 import random
 import time
 from typing import List, AsyncIterator
-from unittest import mock
+import mock
 
 import proto
 import pytest
-from google.auth.aio.transport import Response
 
 from google.api_core import rest_streaming_async
 from google.api import http_pb2
 from google.api import httpbody_pb2
+from google.auth.aio.transport import Response
 
-# TODO (ohmayr): confirm if relative path is not an issue.
-from .._helpers import Composer, Song, EchoResponse, parse_responses
+from ..conftest import Composer, Song, EchoResponse, parse_responses
 
 # TODO (ohmayr): check if we need to log.
 __protobuf__ = proto.module(package=__name__)
@@ -37,9 +36,10 @@
 random.seed(SEED)
 
 
-class AIOHTTPContentMock:
-    def __init__(self, iter_chunked):
-        self.iter_chunked = iter_chunked
+async def mock_async_gen(data, chunk_size=1):
+    for i in range(0, len(data)):
+        chunk = data[i : i + chunk_size]
+        yield chunk.encode("utf-8")
 
 
 class ResponseMock(Response):
@@ -49,6 +49,9 @@ def __init__(self, _response_bytes: bytes, random_split=False):
             self._i = 0
             self._random_split = random_split
 
+        def __aiter__(self):
+            return self
+
         async def __anext__(self):
             if self._i == len(self._responses_bytes):
                 raise StopAsyncIteration
@@ -74,13 +77,15 @@ def __init__(
     async def close(self):
         raise NotImplementedError()
 
-    @property
-    def content(self):
-        iter_chunked = lambda chunk_size: self._ResponseItr(
-            self._parse_responses(),
-            random_split=self._random_split,
+    async def content(self, chunk_size=None):
+        itr = self._ResponseItr(
+            self._parse_responses(), random_split=self._random_split
         )
-        return AIOHTTPContentMock(iter_chunked)
+        async for chunk in itr:
+            yield chunk
+
+    async def read(self):
+        raise NotImplementedError()
 
     @property
     async def headers(self):
@@ -273,14 +278,10 @@ async def test_next_escaped_characters_in_string(
 @pytest.mark.parametrize("response_type", [EchoResponse, httpbody_pb2.HttpBody])
 async def test_next_not_array(response_type):
 
-    mock_content = mock.Mock()
     data = '{"hello": 0}'
-    mock_content.iter_chunked = lambda chunk_size: async_iter(data, chunk_size)
-
     with mock.patch.object(
-        ResponseMock, "content", new_callable=mock.PropertyMock
+        ResponseMock, "content", return_value=mock_async_gen(data)
     ) as mock_method:
-        mock_method.return_value = mock_content
         resp = ResponseMock(responses=[], response_cls=response_type)
         itr = rest_streaming_async.AsyncResponseIterator(resp, response_type)
         with pytest.raises(ValueError):
@@ -321,24 +322,14 @@ async def test_check_buffer(response_type, return_value):
             await itr.__anext__()
 
 
-async def async_iter(data, chunk_size):
-    for i in range(0, len(data) + chunk_size):
-        chunk = data[i : i + chunk_size]
-        yield chunk.encode("utf-8")
-
-
 @pytest.mark.asyncio
 @pytest.mark.parametrize("response_type", [EchoResponse, httpbody_pb2.HttpBody])
 async def test_next_html(response_type):
 
-    mock_content = mock.Mock()
     data = "<!DOCTYPE html><html></html>"
-    mock_content.iter_chunked = lambda chunk_size: async_iter(data, chunk_size)
-
     with mock.patch.object(
-        ResponseMock, "content", new_callable=mock.PropertyMock
+        ResponseMock, "content", return_value=mock_async_gen(data)
     ) as mock_method:
-        mock_method.return_value = mock_content
         resp = ResponseMock(responses=[], response_cls=response_type)
 
         itr = rest_streaming_async.AsyncResponseIterator(resp, response_type)