diff --git a/gcloud/pubsub/message.py b/gcloud/pubsub/message.py index e8db230aa7b6..832a5d554c44 100644 --- a/gcloud/pubsub/message.py +++ b/gcloud/pubsub/message.py @@ -71,6 +71,6 @@ def from_api_repr(cls, api_repr): :type api_repr: dict or None :param api_repr: The API representation of the message """ - data = base64.b64decode(api_repr['data']) + data = base64.b64decode(api_repr.get('data', b'')) return cls(data=data, message_id=api_repr['messageId'], attributes=api_repr.get('attributes')) diff --git a/gcloud/pubsub/test_message.py b/gcloud/pubsub/test_message.py index 693f34405289..7bdc245ea2a9 100644 --- a/gcloud/pubsub/test_message.py +++ b/gcloud/pubsub/test_message.py @@ -42,31 +42,6 @@ def test_ctor_w_attributes(self): self.assertEqual(message.message_id, MESSAGE_ID) self.assertEqual(message.attributes, ATTRS) - def test_from_api_repr_no_attributes(self): - from base64 import b64encode as b64 - DATA = b'DEADBEEF' - B64_DATA = b64(DATA) - MESSAGE_ID = '12345' - api_repr = {'data': B64_DATA, 'messageId': MESSAGE_ID} - message = self._getTargetClass().from_api_repr(api_repr) - self.assertEqual(message.data, DATA) - self.assertEqual(message.message_id, MESSAGE_ID) - self.assertEqual(message.attributes, {}) - - def test_from_api_repr_w_attributes(self): - from base64 import b64encode as b64 - DATA = b'DEADBEEF' - B64_DATA = b64(DATA) - MESSAGE_ID = '12345' - ATTRS = {'a': 'b'} - api_repr = {'data': B64_DATA, - 'messageId': MESSAGE_ID, - 'attributes': ATTRS} - message = self._getTargetClass().from_api_repr(api_repr) - self.assertEqual(message.data, DATA) - self.assertEqual(message.message_id, MESSAGE_ID) - self.assertEqual(message.attributes, ATTRS) - def test_timestamp_no_attributes(self): DATA = b'DEADBEEF' MESSAGE_ID = b'12345' @@ -102,3 +77,36 @@ def test_timestamp_w_timestamp_in_attributes(self): message = self._makeOne(data=DATA, message_id=MESSAGE_ID, attributes=ATTRS) self.assertEqual(message.timestamp, timestamp) + + def test_from_api_repr_missing_data(self): + MESSAGE_ID = '12345' + api_repr = {'messageId': MESSAGE_ID} + message = self._getTargetClass().from_api_repr(api_repr) + self.assertEqual(message.data, b'') + self.assertEqual(message.message_id, MESSAGE_ID) + self.assertEqual(message.attributes, {}) + + def test_from_api_repr_no_attributes(self): + from base64 import b64encode as b64 + DATA = b'DEADBEEF' + B64_DATA = b64(DATA) + MESSAGE_ID = '12345' + api_repr = {'data': B64_DATA, 'messageId': MESSAGE_ID} + message = self._getTargetClass().from_api_repr(api_repr) + self.assertEqual(message.data, DATA) + self.assertEqual(message.message_id, MESSAGE_ID) + self.assertEqual(message.attributes, {}) + + def test_from_api_repr_w_attributes(self): + from base64 import b64encode as b64 + DATA = b'DEADBEEF' + B64_DATA = b64(DATA) + MESSAGE_ID = '12345' + ATTRS = {'a': 'b'} + api_repr = {'data': B64_DATA, + 'messageId': MESSAGE_ID, + 'attributes': ATTRS} + message = self._getTargetClass().from_api_repr(api_repr) + self.assertEqual(message.data, DATA) + self.assertEqual(message.message_id, MESSAGE_ID) + self.assertEqual(message.attributes, ATTRS)