feat: assimilate the gcm client code

Closes: #1057

autopush/router/gcm.py

@@ -1,13 +1,13 @@
 """GCM Router"""
 from typing import Any  # noqa
 
-import gcmclient
 from requests.exceptions import ConnectionError
 from twisted.internet.threads import deferToThread
 from twisted.logger import Logger
 
 from autopush.exceptions import RouterException
 from autopush.metrics import make_tags
+from autopush.router import gcmclient
 from autopush.router.interface import RouterResponse
 from autopush.types import JSONDict  # noqa
 
@@ -115,19 +115,24 @@ def _route(self, notification, uaid_data):
         try:
             gcm = self.gcm[router_data['creds']['senderID']]
             result = gcm.send(payload)
+        except RouterException:
+            raise  # pragma nocover
         except KeyError:
             self.log.critical("Missing GCM bridge credentials")
-            raise RouterException("Server error", status_code=500)
+            raise RouterException("Server error", status_code=500,
+                                  errno=900)
         except gcmclient.GCMAuthenticationError as e:
             self.log.error("GCM Authentication Error: %s" % e)
-            raise RouterException("Server error", status_code=500)
+            raise RouterException("Server error", status_code=500,
+                                  errno=901)
         except ConnectionError as e:
             self.log.warn("GCM Unavailable: %s" % e)
             self.metrics.increment("notification.bridge.error",
                                    tags=make_tags(
                                        self._base_tags,
                                        reason="connection_unavailable"))
             raise RouterException("Server error", status_code=502,
+                                  errno=902,
                                   log_exception=False)
         except Exception as e:
             self.log.error("Unhandled exception in GCM Routing: %s" % e)
@@ -181,15 +186,19 @@ def _process_reply(self, reply, uaid_data, ttl, notification):
                                   )
 
         # retries:
-        if reply.needs_retry():
+        if reply.retry_after:
             self.metrics.increment("notification.bridge.error",
                                    tags=make_tags(self._base_tags,
                                                   reason="retry"))
             self.log.warn("GCM retry requested: {failed()}",
                           failed=lambda: repr(reply.failed.items()))
             raise RouterException("GCM failure to deliver, retry",
                                   status_code=503,
-                                  response_body="Please try request later.",
+                                  headers={"Retry-After": reply.retry_after},
+                                  response_body="Please try request "
+                                                "in {} seconds.".format(
+                                       reply.retry_after
+                                  ),
                                   log_exception=False)
 
         self.metrics.increment("notification.bridge.sent",

autopush/router/gcmclient.py

@@ -0,0 +1,110 @@
+import json
+
+import requests
+
+from autopush.exceptions import RouterException
+
+
+class GCMAuthenticationError(Exception):
+    pass
+
+
+class Result(object):
+
+    def __init__(self, message, response):
+        self.success = {}
+        self.canonicals = {}
+        self.unavailable = []
+        self.not_registered = []
+        self.failed = {}
+
+        self.message = message
+        self.retry_message = None
+
+        self.retry_after = response.headers.get('Retry-After', None)
+
+        if response.status_code != 200:
+            self.retry_message = message
+        else:
+            self._parse_response(message, response.content)
+
+    def _parse_response(self, message, content):
+        data = json.loads(content)
+        if not data.get('results'):
+            raise RouterException("Recv'd invalid response from GCM")
+        reg_id = message.payload['registration_ids'][0]
+        for res in data['results']:
+            if 'message_id' in res:
+                self.success[reg_id] = res['message_id']
+                if 'registration_id' in res:
+                    self.canonicals[reg_id] = res['registration_id']
+            else:
+                if res['error'] in ['Unavailable', 'InternalServerError']:
+                    self.unavailable.append(reg_id)
+                elif res['error'] == 'NotRegistered':
+                    self.not_registered.append(reg_id)
+                else:
+                    self.failed[reg_id] = res['error']
+
+
+class JSONMessage(object):
+
+    def __init__(self,
+                 registration_ids,
+                 collapse_key,
+                 time_to_live,
+                 dry_run,
+                 data):
+        if not registration_ids:
+            raise RouterException("No Registration IDs specified")
+        if not isinstance(registration_ids, list):
+            registration_ids = [registration_ids]
+        self.registration_ids = registration_ids
+        self.payload = {
+            'registration_ids': self.registration_ids,
+            'collapse_key': collapse_key,
+            'time_to_live': int(time_to_live),
+            'delay_while_idle': False,
+            'dry_run': bool(dry_run),
+        }
+        if data:
+            self.payload['data'] = data
+
+
+class GCM(object):
+
+    def __init__(self,
+                 api_key=None,
+                 logger=None,
+                 metrics=None,
+                 endpoint="gcm-http.googleapis.com/gcm/send",
+                 **options):
+
+        self._endpoint = "https://{}".format(endpoint)
+        self._api_key = api_key
+        self.metrics = metrics
+        self.log = logger
+        self._options = options
+        self._sender = requests.post
+
+    def send(self, payload):
+        headers = {
+            'Content-Type': 'application/json',
+            'Authorization': 'key={}'.format(self._api_key),
+        }
+
+        response = self._sender(
+            url=self._endpoint,
+            headers=headers,
+            data=json.dumps(payload.payload),
+            **self._options
+        )
+
+        if response.status_code in (400, 404):
+            raise RouterException(response.content)
+
+        if response.status_code == 401:
+            raise GCMAuthenticationError("Authentication Error")
+
+        if response.status_code == 200 or (500 <= response.status_code <= 599):
+            return Result(payload, response)

autopush/tests/test_gcmclient.py

@@ -0,0 +1,177 @@
+import json
+
+import pytest
+import requests
+from mock import Mock
+from twisted.trial import unittest
+
+from autopush.exceptions import RouterException
+from autopush.router import gcmclient
+
+
+class GCMClientTestCase(unittest.TestCase):
+
+    def setUp(self):
+        self.gcm = gcmclient.GCM(api_key="FakeValue")
+        self.gcm._sender = self._m_request = Mock(spec=requests.post)
+        self._m_response = Mock(spec=requests.Response)
+        self._m_response.return_value = 200
+        self._m_response.headers = dict()
+        self.m_payload = gcmclient.JSONMessage(
+            registration_ids="some_reg_id",
+            collapse_key="coll_key",
+            time_to_live=60,
+            dry_run=False,
+            data={"foo": "bar"}
+        )
+
+    def test_send(self):
+        self._m_response.status_code = 200
+        self._m_response.content = json.dumps({
+            "multicast_id": 5174939174563864884,
+            "success": 1,
+            "failure": 0,
+            "canonical_ids": 0,
+            "results": [
+                {
+                    "message_id": "0:1510011451922224%7a0e7efbaab8b7cc"
+                }
+            ]
+        })
+        self._m_request.return_value = self._m_response
+        result = self.gcm.send(self.m_payload)
+        assert len(result.failed) == 0
+        assert len(result.canonicals) == 0
+        assert (len(result.success) == 1
+                and self.m_payload.registration_ids[0] in result.success)
+
+    def test_canonical(self):
+        self._m_response.status_code = 200
+        self._m_response.content = json.dumps({
+            "multicast_id": 5174939174563864884,
+            "success": 1,
+            "failure": 0,
+            "canonical_ids": 0,
+            "results": [
+                {
+                    "message_id": "0:1510011451922224%7a0e7efbaab8b7cc",
+                    "registration_id": "otherId",
+                }
+            ]
+        })
+        self._m_request.return_value = self._m_response
+        result = self.gcm.send(self.m_payload)
+        assert len(result.failed) == 0
+        assert len(result.canonicals) == 1
+        assert (len(result.success) == 1
+                and self.m_payload.registration_ids[0] in result.success)
+
+    def test_bad_jsonmessage(self):
+        with pytest.raises(RouterException):
+            self.m_payload = gcmclient.JSONMessage(
+                registration_ids=None,
+                collapse_key="coll_key",
+                time_to_live=60,
+                dry_run=False,
+                data={"foo": "bar"}
+            )
+
+    def test_fail_invalid(self):
+        self._m_response.status_code = 200
+        self._m_response.content = json.dumps({
+            "multicast_id": 5174939174563864884,
+            "success": 0,
+            "failure": 1,
+            "canonical_ids": 0,
+            "results": [
+                {
+                    "error": "InvalidRegistration"
+                }
+            ]
+        })
+        self._m_request.return_value = self._m_response
+        result = self.gcm.send(self.m_payload)
+        assert len(result.failed) == 1
+        assert len(result.success) == 0
+
+    def test_fail_unavailable(self):
+        self._m_response.status_code = 200
+        self._m_response.content = json.dumps({
+            "multicast_id": 5174939174563864884,
+            "success": 0,
+            "failure": 1,
+            "canonical_ids": 0,
+            "results": [
+                {
+                    "error": "Unavailable"
+                }
+            ]
+        })
+        self._m_request.return_value = self._m_response
+        result = self.gcm.send(self.m_payload)
+        assert len(result.unavailable) == 1
+        assert len(result.success) == 0
+
+    def test_fail_not_registered(self):
+        self._m_response.status_code = 200
+        self._m_response.content = json.dumps({
+            "multicast_id": 5174939174563864884,
+            "success": 0,
+            "failure": 1,
+            "canonical_ids": 0,
+            "results": [
+                {
+                    "error": "NotRegistered"
+                }
+            ]
+        })
+        self._m_request.return_value = self._m_response
+        result = self.gcm.send(self.m_payload)
+        assert len(result.not_registered) == 1
+        assert len(result.success) == 0
+
+    def test_fail_bad_response(self):
+        self._m_response.status_code = 200
+        self._m_response.content = json.dumps({
+            "multicast_id": 5174939174563864884,
+            "success": 0,
+            "failure": 1,
+            "canonical_ids": 0,
+        })
+        self._m_request.return_value = self._m_response
+        with pytest.raises(RouterException):
+            self.gcm.send(self.m_payload)
+
+    def test_fail_400(self):
+        self._m_response.status_code = 400
+        self._m_response.content = msg = "Invalid JSON"
+        self._m_request.return_value = self._m_response
+        with pytest.raises(RouterException) as ex:
+            self.gcm.send(self.m_payload)
+        assert ex.value.status_code == 500
+        assert ex.value.message == msg
+
+    def test_fail_404(self):
+        self._m_response.status_code = 404
+        self._m_response.content = msg = "Invalid URL"
+        self._m_request.return_value = self._m_response
+        with pytest.raises(RouterException) as ex:
+            self.gcm.send(self.m_payload)
+        assert ex.value.status_code == 500
+        assert ex.value.message == msg
+
+    def test_fail_401(self):
+        self._m_response.status_code = 401
+        self._m_response.content = "Unauthorized"
+        self._m_request.return_value = self._m_response
+        with pytest.raises(gcmclient.GCMAuthenticationError):
+            self.gcm.send(self.m_payload)
+
+    def test_fail_500(self):
+        self._m_response.status_code = 500
+        self._m_response.content = "OMG"
+        self._m_response.headers['Retry-After'] = 123
+        self._m_request.return_value = self._m_response
+        result = self.gcm.send(self.m_payload)
+        assert 'some_reg_id' in result.retry_message.registration_ids
+        assert result.retry_after == 123

autopush/tests/test_integration.py

@@ -1566,11 +1566,8 @@ class MockReply(object):
         failed_items = dict()
         not_registered = dict()
         failed = dict()
-        _needs_retry = False
-
-        @classmethod
-        def needs_retry(cls=None):
-            return False
+        retry_after = False
+        retry_message = None
 
     def _add_router(self):
         from autopush.router.gcm import GCMRouter
@@ -1635,7 +1632,7 @@ def test_registration(self):
             body=data
         )
 
-        ca_data = self._mock_send.call_args[0][0].data
+        ca_data = self._mock_send.call_args[0][0].payload['data']
         assert response.code == 201
         # ChannelID here MUST match what we got from the registration call.
         # Currently, this is a lowercase, hex UUID without dashes.
@@ -1694,7 +1691,7 @@ def test_registration_aes128gcm(self):
             body=data
         )
 
-        ca_data = self._mock_send.call_args[0][0].data
+        ca_data = self._mock_send.call_args[0][0].payload['data']
         assert response.code == 201
         # ChannelID here MUST match what we got from the registration call.
         # Currently, this is a lowercase, hex UUID without dashes.