From a9f0cab73a02c9c496ab968fa97f24592c9af4fd Mon Sep 17 00:00:00 2001
From: Hiranya Jayathilaka <hiranya911@gmail.com>
Date: Wed, 2 Dec 2020 14:22:53 -0800
Subject: [PATCH] fix(fcm): Converting unexpected gapic runtime errors to
 FirebaseError

---
 firebase_admin/messaging.py |  3 +--
 tests/test_messaging.py     | 53 ++++++++++++++++++++++++++++++++-----
 2 files changed, 47 insertions(+), 9 deletions(-)

diff --git a/firebase_admin/messaging.py b/firebase_admin/messaging.py
index 217cf0a56..7c92a3d8d 100644
--- a/firebase_admin/messaging.py
+++ b/firebase_admin/messaging.py
@@ -16,7 +16,6 @@
 
 import json
 
-import googleapiclient
 from googleapiclient import http
 from googleapiclient import _auth
 import requests
@@ -388,7 +387,7 @@ def batch_callback(_, response, error):
 
         try:
             batch.execute()
-        except googleapiclient.http.HttpError as error:
+        except Exception as error:
             raise self._handle_batch_error(error)
         else:
             return BatchResponse(responses)
diff --git a/tests/test_messaging.py b/tests/test_messaging.py
index 6333aad46..8eb24c0a9 100644
--- a/tests/test_messaging.py
+++ b/tests/test_messaging.py
@@ -1792,6 +1792,15 @@ def test_send_unknown_fcm_error_code(self, status):
         assert json.loads(recorder[0].body.decode()) == body
 
 
+class _HttpMockException:
+
+    def __init__(self, exc):
+        self._exc = exc
+
+    def request(self, url, **kwargs):
+        raise self._exc
+
+
 class TestBatch:
 
     @classmethod
@@ -1803,17 +1812,21 @@ def setup_class(cls):
     def teardown_class(cls):
         testutils.cleanup_apps()
 
-    def _instrument_batch_messaging_service(self, app=None, status=200, payload=''):
+    def _instrument_batch_messaging_service(self, app=None, status=200, payload='', exc=None):
         if not app:
             app = firebase_admin.get_app()
+
         fcm_service = messaging._get_messaging_service(app)
-        if status == 200:
-            content_type = 'multipart/mixed; boundary=boundary'
+        if exc:
+            fcm_service._transport = _HttpMockException(exc)
         else:
-            content_type = 'application/json'
-        fcm_service._transport = http.HttpMockSequence([
-            ({'status': str(status), 'content-type': content_type}, payload),
-        ])
+            if status == 200:
+                content_type = 'multipart/mixed; boundary=boundary'
+            else:
+                content_type = 'application/json'
+            fcm_service._transport = http.HttpMockSequence([
+                ({'status': str(status), 'content-type': content_type}, payload),
+            ])
         return fcm_service
 
     def _batch_payload(self, payloads):
@@ -2027,6 +2040,19 @@ def test_send_all_batch_fcm_error_code(self, status):
             messaging.send_all([msg])
         check_exception(excinfo.value, 'test error', status)
 
+    def test_send_all_runtime_exception(self):
+        exc = BrokenPipeError('Test error')
+        _ = self._instrument_batch_messaging_service(exc=exc)
+        msg = messaging.Message(topic='foo')
+
+        with pytest.raises(exceptions.UnknownError) as excinfo:
+            messaging.send_all([msg])
+
+        expected = 'Unknown error while making a remote service call: Test error'
+        assert str(excinfo.value) == expected
+        assert excinfo.value.cause is exc
+        assert excinfo.value.http_response is None
+
 
 class TestSendMulticast(TestBatch):
 
@@ -2204,6 +2230,19 @@ def test_send_multicast_batch_fcm_error_code(self, status):
             messaging.send_multicast(msg)
         check_exception(excinfo.value, 'test error', status)
 
+    def test_send_multicast_runtime_exception(self):
+        exc = BrokenPipeError('Test error')
+        _ = self._instrument_batch_messaging_service(exc=exc)
+        msg = messaging.MulticastMessage(tokens=['foo'])
+
+        with pytest.raises(exceptions.UnknownError) as excinfo:
+            messaging.send_multicast(msg)
+
+        expected = 'Unknown error while making a remote service call: Test error'
+        assert str(excinfo.value) == expected
+        assert excinfo.value.cause is exc
+        assert excinfo.value.http_response is None
+
 
 class TestTopicManagement: