add context manager methods

google/api_core/rest_streaming_async.py

@@ -20,7 +20,7 @@
 
 try:
     import google.auth.aio.transport
-except ImportError:
+except ImportError:  # pragma: NO COVER
     # TODO (ohmayr): Update this version once auth work is released.
     raise ValueError(
         "google-auth>=2.3x.x is required to use asynchronous rest streaming."
@@ -56,6 +56,9 @@ def __init__(
             response_message_cls=response_message_cls
         )
 
+    async def __aenter__(self):
+        return self
+
     async def cancel(self):
         """Cancel existing streaming operation."""
         await self._response.close()

tests/asyncio/test_rest_streaming_async.py

@@ -54,7 +54,7 @@
 
 
 async def mock_async_gen(data, chunk_size=1):
-    for i in range(0, len(data)):
+    for i in range(0, len(data)):  # pragma: NO COVER
         chunk = data[i : i + chunk_size]
         yield chunk.encode("utf-8")
 
@@ -316,6 +316,18 @@ async def test_cancel(response_type):
         mock_method.assert_called_once()
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("response_type", [EchoResponse, httpbody_pb2.HttpBody])
+async def test_iterator_as_context_manager(response_type):
+    with mock.patch.object(
+        ResponseMock, "close", new_callable=mock.AsyncMock
+    ) as mock_method:
+        resp = ResponseMock(responses=[], response_cls=response_type)
+        async with rest_streaming_async.AsyncResponseIterator(resp, response_type):
+            pass
+        mock_method.assert_called_once()
+
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize(
     "response_type,return_value",