diff --git a/google/api_core/rest_streaming.py b/google/api_core/rest_streaming.py index f91381c1..3f5b6b03 100644 --- a/google/api_core/rest_streaming.py +++ b/google/api_core/rest_streaming.py @@ -16,9 +16,12 @@ from collections import deque import string -from typing import Deque +from typing import Deque, Union +import proto import requests +import google.protobuf.message +from google.protobuf.json_format import Parse class ResponseIterator: @@ -26,11 +29,18 @@ class ResponseIterator: Args: response (requests.Response): An API response object. - response_message_cls (Callable[proto.Message]): A proto + response_message_cls (Union[proto.Message, google.protobuf.message.Message]): A response class expected to be returned from an API. + + Raises: + ValueError: If `response_message_cls` is not a subclass of `proto.Message` or `google.protobuf.message.Message`. """ - def __init__(self, response: requests.Response, response_message_cls): + def __init__( + self, + response: requests.Response, + response_message_cls: Union[proto.Message, google.protobuf.message.Message], + ): self._response = response self._response_message_cls = response_message_cls # Inner iterator over HTTP response's content. @@ -107,7 +117,14 @@ def __next__(self): def _grab(self): # Add extra quotes to make json.loads happy. - return self._response_message_cls.from_json(self._ready_objs.popleft()) + if issubclass(self._response_message_cls, proto.Message): + return self._response_message_cls.from_json(self._ready_objs.popleft()) + elif issubclass(self._response_message_cls, google.protobuf.message.Message): + return Parse(self._ready_objs.popleft(), self._response_message_cls()) + else: + raise ValueError( + "Response message class must be a subclass of proto.Message or google.protobuf.message.Message." + ) def __iter__(self): return self diff --git a/tests/unit/test_rest_streaming.py b/tests/unit/test_rest_streaming.py index a44c83c0..b532eb1d 100644 --- a/tests/unit/test_rest_streaming.py +++ b/tests/unit/test_rest_streaming.py @@ -24,8 +24,11 @@ import requests from google.api_core import rest_streaming +from google.api import http_pb2 +from google.api import httpbody_pb2 from google.protobuf import duration_pb2 from google.protobuf import timestamp_pb2 +from google.protobuf.json_format import MessageToJson __protobuf__ = proto.module(package=__name__) @@ -98,7 +101,10 @@ def _parse_responses(self, responses: List[proto.Message]) -> bytes: # json.dumps returns a string surrounded with quotes that need to be stripped # in order to be an actual JSON. json_responses = [ - self._response_message_cls.to_json(r).strip('"') for r in responses + self._response_message_cls.to_json(r).strip('"') + if issubclass(self._response_message_cls, proto.Message) + else MessageToJson(r).strip('"') + for r in responses ] logging.info(f"Sending JSON stream: {json_responses}") ret_val = "[{}]".format(",".join(json_responses)) @@ -114,103 +120,220 @@ def iter_content(self, *args, **kwargs): ) -@pytest.mark.parametrize("random_split", [False]) -def test_next_simple(random_split): - responses = [EchoResponse(content="hello world"), EchoResponse(content="yes")] +@pytest.mark.parametrize( + "random_split,resp_message_is_proto_plus", + [(False, True), (False, False)], +) +def test_next_simple(random_split, resp_message_is_proto_plus): + if resp_message_is_proto_plus: + response_type = EchoResponse + responses = [EchoResponse(content="hello world"), EchoResponse(content="yes")] + else: + response_type = httpbody_pb2.HttpBody + responses = [ + httpbody_pb2.HttpBody(content_type="hello world"), + httpbody_pb2.HttpBody(content_type="yes"), + ] + resp = ResponseMock( - responses=responses, random_split=random_split, response_cls=EchoResponse + responses=responses, random_split=random_split, response_cls=response_type ) - itr = rest_streaming.ResponseIterator(resp, EchoResponse) + itr = rest_streaming.ResponseIterator(resp, response_type) assert list(itr) == responses -@pytest.mark.parametrize("random_split", [True, False]) -def test_next_nested(random_split): - responses = [ - Song(title="some song", composer=Composer(given_name="some name")), - Song(title="another song", date_added=datetime.datetime(2021, 12, 17)), - ] +@pytest.mark.parametrize( + "random_split,resp_message_is_proto_plus", + [ + (True, True), + (False, True), + (True, False), + (False, False), + ], +) +def test_next_nested(random_split, resp_message_is_proto_plus): + if resp_message_is_proto_plus: + response_type = Song + responses = [ + Song(title="some song", composer=Composer(given_name="some name")), + Song(title="another song", date_added=datetime.datetime(2021, 12, 17)), + ] + else: + # Although `http_pb2.HttpRule`` is used in the response, any response message + # can be used which meets this criteria for the test of having a nested field. + response_type = http_pb2.HttpRule + responses = [ + http_pb2.HttpRule( + selector="some selector", + custom=http_pb2.CustomHttpPattern(kind="some kind"), + ), + http_pb2.HttpRule( + selector="another selector", + custom=http_pb2.CustomHttpPattern(path="some path"), + ), + ] resp = ResponseMock( - responses=responses, random_split=random_split, response_cls=Song + responses=responses, random_split=random_split, response_cls=response_type ) - itr = rest_streaming.ResponseIterator(resp, Song) + itr = rest_streaming.ResponseIterator(resp, response_type) assert list(itr) == responses -@pytest.mark.parametrize("random_split", [True, False]) -def test_next_stress(random_split): +@pytest.mark.parametrize( + "random_split,resp_message_is_proto_plus", + [ + (True, True), + (False, True), + (True, False), + (False, False), + ], +) +def test_next_stress(random_split, resp_message_is_proto_plus): n = 50 - responses = [ - Song(title="title_%d" % i, composer=Composer(given_name="name_%d" % i)) - for i in range(n) - ] + if resp_message_is_proto_plus: + response_type = Song + responses = [ + Song(title="title_%d" % i, composer=Composer(given_name="name_%d" % i)) + for i in range(n) + ] + else: + response_type = http_pb2.HttpRule + responses = [ + http_pb2.HttpRule( + selector="selector_%d" % i, + custom=http_pb2.CustomHttpPattern(path="path_%d" % i), + ) + for i in range(n) + ] resp = ResponseMock( - responses=responses, random_split=random_split, response_cls=Song + responses=responses, random_split=random_split, response_cls=response_type ) - itr = rest_streaming.ResponseIterator(resp, Song) + itr = rest_streaming.ResponseIterator(resp, response_type) assert list(itr) == responses -@pytest.mark.parametrize("random_split", [True, False]) -def test_next_escaped_characters_in_string(random_split): - composer_with_relateds = Composer() - relateds = ["Artist A", "Artist B"] - composer_with_relateds.relateds = relateds - - responses = [ - Song(title='ti"tle\nfoo\tbar{}', composer=Composer(given_name="name\n\n\n")), - Song( - title='{"this is weird": "totally"}', composer=Composer(given_name="\\{}\\") - ), - Song(title='\\{"key": ["value",]}\\', composer=composer_with_relateds), - ] +@pytest.mark.parametrize( + "random_split,resp_message_is_proto_plus", + [ + (True, True), + (False, True), + (True, False), + (False, False), + ], +) +def test_next_escaped_characters_in_string(random_split, resp_message_is_proto_plus): + if resp_message_is_proto_plus: + response_type = Song + composer_with_relateds = Composer() + relateds = ["Artist A", "Artist B"] + composer_with_relateds.relateds = relateds + + responses = [ + Song( + title='ti"tle\nfoo\tbar{}', composer=Composer(given_name="name\n\n\n") + ), + Song( + title='{"this is weird": "totally"}', + composer=Composer(given_name="\\{}\\"), + ), + Song(title='\\{"key": ["value",]}\\', composer=composer_with_relateds), + ] + else: + response_type = http_pb2.Http + responses = [ + http_pb2.Http( + rules=[ + http_pb2.HttpRule( + selector='ti"tle\nfoo\tbar{}', + custom=http_pb2.CustomHttpPattern(kind="name\n\n\n"), + ) + ] + ), + http_pb2.Http( + rules=[ + http_pb2.HttpRule( + selector='{"this is weird": "totally"}', + custom=http_pb2.CustomHttpPattern(kind="\\{}\\"), + ) + ] + ), + http_pb2.Http( + rules=[ + http_pb2.HttpRule( + selector='\\{"key": ["value",]}\\', + custom=http_pb2.CustomHttpPattern(kind="\\{}\\"), + ) + ] + ), + ] resp = ResponseMock( - responses=responses, random_split=random_split, response_cls=Song + responses=responses, random_split=random_split, response_cls=response_type ) - itr = rest_streaming.ResponseIterator(resp, Song) + itr = rest_streaming.ResponseIterator(resp, response_type) assert list(itr) == responses -def test_next_not_array(): +@pytest.mark.parametrize("response_type", [EchoResponse, httpbody_pb2.HttpBody]) +def test_next_not_array(response_type): with patch.object( ResponseMock, "iter_content", return_value=iter('{"hello": 0}') ) as mock_method: - - resp = ResponseMock(responses=[], response_cls=EchoResponse) - itr = rest_streaming.ResponseIterator(resp, EchoResponse) + resp = ResponseMock(responses=[], response_cls=response_type) + itr = rest_streaming.ResponseIterator(resp, response_type) with pytest.raises(ValueError): next(itr) mock_method.assert_called_once() -def test_cancel(): +@pytest.mark.parametrize("response_type", [EchoResponse, httpbody_pb2.HttpBody]) +def test_cancel(response_type): with patch.object(ResponseMock, "close", return_value=None) as mock_method: - resp = ResponseMock(responses=[], response_cls=EchoResponse) - itr = rest_streaming.ResponseIterator(resp, EchoResponse) + resp = ResponseMock(responses=[], response_cls=response_type) + itr = rest_streaming.ResponseIterator(resp, response_type) itr.cancel() mock_method.assert_called_once() -def test_check_buffer(): +@pytest.mark.parametrize( + "response_type,return_value", + [ + (EchoResponse, bytes('[{"content": "hello"}, {', "utf-8")), + (httpbody_pb2.HttpBody, bytes('[{"content_type": "hello"}, {', "utf-8")), + ], +) +def test_check_buffer(response_type, return_value): with patch.object( ResponseMock, "_parse_responses", - return_value=bytes('[{"content": "hello"}, {', "utf-8"), + return_value=return_value, ): - resp = ResponseMock(responses=[], response_cls=EchoResponse) - itr = rest_streaming.ResponseIterator(resp, EchoResponse) + resp = ResponseMock(responses=[], response_cls=response_type) + itr = rest_streaming.ResponseIterator(resp, response_type) with pytest.raises(ValueError): next(itr) next(itr) -def test_next_html(): +@pytest.mark.parametrize("response_type", [EchoResponse, httpbody_pb2.HttpBody]) +def test_next_html(response_type): with patch.object( ResponseMock, "iter_content", return_value=iter("") ) as mock_method: - - resp = ResponseMock(responses=[], response_cls=EchoResponse) - itr = rest_streaming.ResponseIterator(resp, EchoResponse) + resp = ResponseMock(responses=[], response_cls=response_type) + itr = rest_streaming.ResponseIterator(resp, response_type) with pytest.raises(ValueError): next(itr) mock_method.assert_called_once() + + +def test_invalid_response_class(): + class SomeClass: + pass + + resp = ResponseMock(responses=[], response_cls=SomeClass) + response_iterator = rest_streaming.ResponseIterator(resp, SomeClass) + with pytest.raises( + ValueError, + match="Response message class must be a subclass of proto.Message or google.protobuf.message.Message", + ): + response_iterator._grab()