Added support for timeouts in messaging module

carsongee · carsongee · commit 0f6cf00bf99c · 2018-03-08T19:07:22.000-05:00
diff --git a/.gitignore b/.gitignore
@@ -9,3 +9,5 @@ dist/
 *~
 scripts/cert.json
 scripts/apikey.txt
+htmlcov/
+.pytest_cache/
diff --git a/firebase_admin/messaging.py b/firebase_admin/messaging.py
@@ -748,6 +748,7 @@ def __init__(self, app):
                 'GCLOUD_PROJECT environment variable.')
         self._fcm_url = _MessagingService.FCM_URL.format(project_id)
         self._client = _http_client.JsonHttpClient(credential=app.credential.get_credential())
+        self._timeout = app.options.get('httpTimeout')
 
     @classmethod
     def encode_message(cls, message):
@@ -760,7 +761,7 @@ def send(self, message, dry_run=False):
         if dry_run:
             data['validate_only'] = True
         try:
-            resp = self._client.body('post', url=self._fcm_url, json=data)
+            resp = self._client.body('post', url=self._fcm_url, json=data, timeout=self._timeout)
         except requests.exceptions.RequestException as error:
             if error.response is not None:
                 self._handle_fcm_error(error)
@@ -791,7 +792,12 @@ def make_topic_management_request(self, tokens, topic, operation):
         url = '{0}/{1}'.format(_MessagingService.IID_URL, operation)
         try:
             resp = self._client.body(
-                'post', url=url, json=data, headers=_MessagingService.IID_HEADERS)
+                'post',
+                url=url,
+                json=data,
+                headers=_MessagingService.IID_HEADERS,
+                timeout=self._timeout
+            )
         except requests.exceptions.RequestException as error:
             if error.response is not None:
                 self._handle_iid_error(error)
diff --git a/tests/test_messaging.py b/tests/test_messaging.py
@@ -750,6 +750,44 @@ def test_aps_alert(self):
         check_encoding(msg, expected)
 
 
+class TestTimeout(object):
+
+    @classmethod
+    def setup_class(cls):
+        cred = testutils.MockCredential()
+        firebase_admin.initialize_app(cred, {'httpTimeout': 4, 'projectId': 'explicit-project-id'})
+
+    @classmethod
+    def teardown_class(cls):
+        testutils.cleanup_apps()
+
+    def setup(self):
+        app = firebase_admin.get_app()
+        self.fcm_service = messaging._get_messaging_service(app)
+        self.recorder = []
+
+    def test_send(self):
+        self.fcm_service._client.session.mount(
+            'https://fcm.googleapis.com',
+            testutils.MockAdapter(json.dumps({'name': 'message-id'}), 200, self.recorder))
+        msg = messaging.Message(topic='foo')
+        messaging.send(msg)
+        assert len(self.recorder) == 1
+        assert self.recorder[0]._extra_kwargs['timeout'] == 4
+
+    def test_topic_management_timeout(self):
+        self.fcm_service._client.session.mount(
+            'https://iid.googleapis.com',
+            testutils.MockAdapter(
+                json.dumps({'results': [{}, {'error': 'error_reason'}]}),
+                200,
+                self.recorder)
+        )
+        messaging.subscribe_to_topic(['1'], 'a')
+        assert len(self.recorder) == 1
+        assert self.recorder[0]._extra_kwargs['timeout'] == 4
+
+
 class TestSend(object):
 
     _DEFAULT_RESPONSE = json.dumps({'name': 'message-id'})
@@ -817,6 +855,7 @@ def test_send(self):
         assert len(recorder) == 1
         assert recorder[0].method == 'POST'
         assert recorder[0].url == self._get_url('explicit-project-id')
+        assert recorder[0]._extra_kwargs['timeout'] == None
         body = {'message': messaging._MessagingService.encode_message(msg)}
         assert json.loads(recorder[0].body.decode()) == body
 
@@ -875,7 +914,6 @@ def test_send_canonical_error_code(self, status):
         body = {'message': messaging._MessagingService.JSON_ENCODER.default(msg)}
         assert json.loads(recorder[0].body.decode()) == body
 
-
 class TestTopicManagement(object):
 
     _DEFAULT_RESPONSE = json.dumps({'results': [{}, {'error': 'error_reason'}]})
