Merge pull request #1481 from tseaver/1480-pubsub-empty_payload-missi…

…ng_data Harden 'Message.from_api_repr' against missing 'data'
googleapis · Feb 17, 2016 · 87c7a4c · 87c7a4c
2 parents 75a395f + 3b9b0b6
commit 87c7a4c
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 26 deletions.
diff --git 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
@@ -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)