Move batch message encoding in to _messaging_utils class BatchMessageEncoder

Author: ZachOrr

firebase_admin/_messaging_utils.py

@@ -773,3 +773,35 @@ def default(self, obj): # pylint: disable=method-hidden
         if target_count != 1:
             raise ValueError('Exactly one of token, topic or condition must be specified.')
         return result
+
+
+class BatchMessageEncoder(object):
+    """A custom encoder for serializing JSON for Messages into batch request data."""
+
+    @classmethod
+    def encode(cls, url, headers, boundary, message_data): # pylint: disable=method-hidden
+        parts = [cls._batch_request_part(url, headers, boundary, md, index) for (index, md) in enumerate(message_data)]
+        return '{}--{}--\r\n'.format(''.join(parts), boundary)
+
+    @classmethod
+    def _batch_request_part(cls, url, headers, boundary, message_data, index):
+        data = cls._batch_request_part_data(url, headers, message_data)
+        part = '--{}\r\n'.format(boundary)
+        part += 'Content-Length: {}\r\n'.format(len(data))
+        part += 'Content-Type: application/http\r\n'
+        part += 'content-id: {}\r\n'.format(index + 1)
+        part += 'content-transfer-encoding: binary\r\n'
+        part += '\r\n'
+        part += '{}\r\n'.format(data)
+        return part
+
+    @classmethod
+    def _batch_request_part_data(cls, url, headers, message_data):
+        body = json.dumps(message_data)
+        data = 'POST {} HTTP/1.1\r\n'.format(url)
+        data += 'Content-Length: {}\r\n'.format(len(body))
+        data += 'Content-Type: application/json; charset=UTF-8\r\n'
+        data += '\r\n'.join(['{}: {}'.format(k, v) for k, v in headers.items()])
+        data += '\r\n\r\n'
+        data += body
+        return data

firebase_admin/messaging.py

@@ -252,6 +252,7 @@ class _MessagingService(object):
     IID_URL = 'https://iid.googleapis.com'
     IID_HEADERS = {'access_token_auth': 'true'}
     JSON_ENCODER = _messaging_utils.MessageEncoder()
+    BATCH_ENCODER = _messaging_utils.BatchMessageEncoder()
 
     BATCH_PART_BOUNDARY = '__END_OF_PART__'
 
@@ -317,7 +318,13 @@ def send(self, message, dry_run=False):
             return resp['name']
 
     def sendAll(self, messages, dry_run=False):
-        data = self._batch_message_data(messages, dry_run)
+        message_data = [self._message_data(message, dry_run) for message in messages]
+        data = _MessagingService.BATCH_ENCODER.encode(
+            url=self._fcm_url,
+            headers=self._fcm_headers,
+            boundary=_MessagingService.BATCH_PART_BOUNDARY,
+            message_data=message_data
+        )
         headers = {
             'Content-Type': 'multipart/mixed; boundary={}'.format(_MessagingService.BATCH_PART_BOUNDARY)
         }
@@ -386,33 +393,6 @@ def _message_data(self, message, dry_run):
             data['validate_only'] = True
         return data
 
-    def _batch_message_data(self, messages, dry_run):
-        parts = [self._batch_request_part(request, dry_run, _MessagingService.BATCH_PART_BOUNDARY, index)\
-            for (index, request) in enumerate(messages)]
-        return '{}--{}--\r\n'.format(''.join(parts), _MessagingService.BATCH_PART_BOUNDARY)
-
-    def _batch_request_part(self, request, dry_run, boundary, index):
-        data = self._batch_request_part_data(request, dry_run)
-        part = '--{}\r\n'.format(boundary)
-        part += 'Content-Length: {}\r\n'.format(len(data))
-        part += 'Content-Type: application/http\r\n'
-        part += 'content-id: {}\r\n'.format(index + 1)
-        part += 'content-transfer-encoding: binary\r\n'
-        part += '\r\n'
-        part += '{}\r\n'.format(data)
-        return part
-
-    def _batch_request_part_data(self, request, dry_run):
-        body = json.dumps(self._message_data(request, dry_run))
-        data = 'POST {} HTTP/1.1\r\n'.format(self._fcm_url)
-        data += 'Content-Length: {}\r\n'.format(len(body))
-        data += 'Content-Type: application/json; charset=UTF-8\r\n'
-        data += '\r\n'.join(['{}: {}'.format(k, v) for k, v in self._fcm_headers.items()])
-        data += '\r\n\r\n'
-        data += body
-        return data
-
-
     def _handle_fcm_error(self, error):
         """Handles errors received from the FCM API."""
         data = {}