Merge pull request #6 from jrconlin/bug/04

feat: Add tests for GCM compatibility

Author: jrconlin

pywebpush/__init__.py

@@ -41,6 +41,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
@@ -173,7 +177,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
@@ -183,7 +186,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('='),
                 'p256dh': base64.urlsafe_b64encode(
@@ -114,3 +115,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")