feat: Add tests for GCM compatibility

Fixed some Python3 GCM issues

Author: jrconlin

pywebpush/__init__.py

@@ -3,10 +3,10 @@
 # file, You can obtain one at http://mozilla.org/MPL/2.0/.
 
 import base64
-import os
-
 import json
+import os
 
+import six
 import http_ece
 import pyelliptic
 import requests
@@ -42,6 +42,10 @@ def get(self, key, default=None):
         except KeyError:
             return default
 
+    def update(self, data):
+        for key in data:
+            self.__setitem__(key, data[key])
+
 
 class WebPusher:
     """WebPusher encrypts a data block using HTTP Encrypted Content Encoding
@@ -88,11 +92,6 @@ def __init__(self, subscription_info):
             the client.
 
         """
-        # Python 2 v. 3 hack
-        try:
-            self.basestr = basestring
-        except NameError:
-            self.basestr = str
         if 'endpoint' not in subscription_info:
             raise WebPushException("subscription_info missing endpoint URL")
         if 'keys' not in subscription_info:
@@ -102,7 +101,7 @@ def __init__(self, subscription_info):
         for k in ['p256dh', 'auth']:
             if keys.get(k) is None:
                 raise WebPushException("Missing keys value: %s", k)
-            if isinstance(keys[k], self.basestr):
+            if isinstance(keys[k], six.string_types):
                 keys[k] = bytes(keys[k].encode('utf8'))
         receiver_raw = base64.urlsafe_b64decode(self._repad(keys['p256dh']))
         if len(receiver_raw) != 65 and receiver_raw[0] != "\x04":
@@ -131,7 +130,7 @@ def encode(self, data):
         # ID tag.
         server_key_id = base64.urlsafe_b64encode(server_key.get_pubkey()[1:])
 
-        if isinstance(data, self.basestr):
+        if isinstance(data, six.string_types):
             data = bytes(data.encode('utf8'))
 
         # http_ece requires that these both be set BEFORE encrypt or
@@ -183,7 +182,6 @@ def send(self, data, headers={}, ttl=0, gcm_key=None, reg_id=None):
         })
         gcm_endpoint = 'https://android.googleapis.com/gcm/send'
         if self.subscription_info['endpoint'].startswith(gcm_endpoint):
-
             if not gcm_key:
                 raise WebPushException("API key not provided for gcm endpoint")
             endpoint = gcm_endpoint
@@ -193,7 +191,8 @@ def send(self, data, headers={}, ttl=0, gcm_key=None, reg_id=None):
             reg_ids.append(reg_id)
             data = {}
             data['registration_ids'] = reg_ids
-            data['raw_data'] = base64.b64encode(encoded.get('body'))
+            data['raw_data'] = base64.b64encode(
+                encoded.get('body')).decode('utf8')
             encoded_data = json.dumps(data)
             headers.update({
                 'Authorization': 'key='+gcm_key,

pywebpush/tests/test_webpush.py

@@ -1,4 +1,5 @@
 import base64
+import json
 import os
 import unittest
 
@@ -11,9 +12,9 @@
 
 
 class WebpushTestCase(unittest.TestCase):
-    def _gen_subscription_info(self, recv_key):
+    def _gen_subscription_info(self, recv_key, endpoint="https://example.com"):
         return {
-            "endpoint": "https://example.com/",
+            "endpoint": endpoint,
             "keys": {
                 'auth': base64.urlsafe_b64encode(os.urandom(16)).strip(b'='),
                 'p256dh': base64.urlsafe_b64encode(
@@ -115,3 +116,25 @@ def test_ci_dict(self):
         eq_('apple', ci.get("Foo"))
         del (ci['FOO'])
         eq_(None, ci.get('Foo'))
+
+    @patch("requests.post")
+    def test_gcm(self, mock_post):
+        recv_key = pyelliptic.ECC(curve="prime256v1")
+        subscription_info = self._gen_subscription_info(
+            recv_key,
+            endpoint="https://android.googleapis.com/gcm/send/regid123")
+        headers = {"Crypto-Key": "pre-existing",
+                   "Authentication": "bearer vapid"}
+        data = "Mary had a little lamb"
+        wp = WebPusher(subscription_info)
+        self.assertRaises(
+            WebPushException,
+            wp.send,
+            data,
+            headers)
+        wp.send(data, headers, gcm_key="gcm_key_value")
+        pdata = json.loads(mock_post.call_args[1].get('data'))
+        pheaders = mock_post.call_args[1].get('headers')
+        eq_(pdata["registration_ids"][0], "regid123")
+        eq_(pheaders.get("authorization"), "key=gcm_key_value")
+        eq_(pheaders.get("content-type"), "application/json")