Fix pubsub errors when receiving unencoded message. googleapis#2513.

daspecster · daspecster · commit 6fac5d1d8a66 · 2016-10-12T10:24:49.000-04:00
diff --git a/docs/pubsub-usage.rst b/docs/pubsub-usage.rst
@@ -85,6 +85,10 @@ Test permissions allowed by the current IAM policy on a topic:
 Publish messages to a topic
 ---------------------------
 
+.. note::
+    If you're publishing a message from console.cloud.google.com, you will need
+    to base64 encode the message first.
+
 Publish a single message to a topic, without attributes:
 
 .. literalinclude:: pubsub_snippets.py
diff --git a/pubsub/google/cloud/pubsub/message.py b/pubsub/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'))
diff --git a/pubsub/unit_tests/test_message.py b/pubsub/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'
