Merge pull request #2527 from daspecster/fix-pubsub-handling-unencoded-data

Fix pubsub errors when receiving unencoded message.

Author: daspecster

google/cloud/pubsub/message.py

@@ -15,6 +15,7 @@
 """Define API Topics."""
 
 import base64
+import binascii
 
 from google.cloud._helpers import _rfc3339_to_datetime
 
@@ -85,7 +86,14 @@ def from_api_repr(cls, api_repr):
         :rtype: :class:`Message`
         :returns: The message created from the response.
         """
-        data = base64.b64decode(api_repr.get('data', b''))
+        raw_data = api_repr.get('data', b'')
+        try:
+            data = base64.b64decode(raw_data)
+        except (binascii.Error, TypeError):
+            to_pad = (- len(raw_data)) % 4
+            padded_data = raw_data + b'=' * to_pad
+            data = base64.b64decode(padded_data)
+
         instance = cls(
             data=data, message_id=api_repr['messageId'],
             attributes=api_repr.get('attributes'))

unit_tests/test_message.py

@@ -89,6 +89,19 @@ def test_from_api_repr_missing_data(self):
         self.assertEqual(message.attributes, {})
         self.assertIsNone(message.service_timestamp)
 
+    def test_from_api_repr_bad_b64_data(self):
+        DATA = b'wefwefw'
+        BAD_B64_DATA = b'd2Vmd2Vmdw='
+        MESSAGE_ID = '12345'
+        TIMESTAMP = '2016-03-18-19:38:22.001393427Z'
+        api_repr = {
+            'data': BAD_B64_DATA,
+            'messageId': MESSAGE_ID,
+            'publishTimestamp': TIMESTAMP,
+        }
+        message = self._getTargetClass().from_api_repr(api_repr)
+        self.assertEqual(message.data, DATA)
+
     def test_from_api_repr_no_attributes(self):
         from base64 import b64encode as b64
         DATA = b'DEADBEEF'