Add support for arbitrary key-value pairs in messaging.ApsAlert

viktorasl · viktorasl · commit 69faac3a17f6 · 2019-08-14T16:16:02.000+03:00
diff --git a/firebase_admin/_messaging_utils.py b/firebase_admin/_messaging_utils.py
@@ -397,7 +397,8 @@ class ApsAlert(object):
     """
 
     def __init__(self, title=None, subtitle=None, body=None, loc_key=None, loc_args=None,
-                 title_loc_key=None, title_loc_args=None, action_loc_key=None, launch_image=None):
+                 title_loc_key=None, title_loc_args=None, action_loc_key=None, launch_image=None,
+                 custom_data=None):
         self.title = title
         self.subtitle = subtitle
         self.body = body
@@ -407,6 +408,7 @@ def __init__(self, title=None, subtitle=None, body=None, loc_key=None, loc_args=
         self.title_loc_args = title_loc_args
         self.action_loc_key = action_loc_key
         self.launch_image = launch_image
+        self.custom_data = custom_data
 
 
 class APNSFcmOptions(object):
@@ -835,6 +837,14 @@ def encode_aps_alert(cls, alert):
         if result.get('title-loc-args') and not result.get('title-loc-key'):
             raise ValueError(
                 'ApsAlert.title_loc_key is required when specifying title_loc_args.')
+        if alert.custom_data is not None:
+            if not isinstance(alert.custom_data, dict):
+                raise ValueError('ApsAlert.custom_data must be a dict.')
+            for key, val in alert.custom_data.items():
+                _Validators.check_string('ApsAlert.custom_data key', key)
+                # allow specifying key override because Apple could update API so that key
+                # could have unexpected value type
+                result[key] = val
         return cls.remove_null_values(result)
 
     @classmethod
diff --git a/tests/test_messaging.py b/tests/test_messaging.py
@@ -1209,6 +1209,72 @@ def test_aps_alert(self):
         }
         check_encoding(msg, expected)
 
+    def test_aps_alert_custom_data_merge(self):
+        msg = messaging.Message(
+            topic='topic',
+            apns=messaging.APNSConfig(
+                payload=messaging.APNSPayload(
+                    aps=messaging.Aps(
+                        alert=messaging.ApsAlert(
+                            title='t',
+                            subtitle='st',
+                            custom_data={'k1': 'v1', 'k2': 'v2'}
+                        )
+                    ),
+                )
+            )
+        )
+        expected = {
+            'topic': 'topic',
+            'apns': {
+                'payload': {
+                    'aps': {
+                        'alert': {
+                            'title': 't',
+                            'subtitle': 'st',
+                            'k1': 'v1',
+                            'k2': 'v2'
+                        },
+                    },
+                }
+            },
+        }
+        check_encoding(msg, expected)
+
+    def test_aps_alert_custom_data_override(self):
+        msg = messaging.Message(
+            topic='topic',
+            apns=messaging.APNSConfig(
+                payload=messaging.APNSPayload(
+                    aps=messaging.Aps(
+                        alert=messaging.ApsAlert(
+                            title='t',
+                            subtitle='st',
+                            launch_image='li',
+                            custom_data={'launch-image': ['li1', 'li2']}
+                        )
+                    ),
+                )
+            )
+        )
+        expected = {
+            'topic': 'topic',
+            'apns': {
+                'payload': {
+                    'aps': {
+                        'alert': {
+                            'title': 't',
+                            'subtitle': 'st',
+                            'launch-image': [
+                                'li1',
+                                'li2'
+                            ]
+                        },
+                    },
+                }
+            },
+        }
+        check_encoding(msg, expected)
 
 class TestTimeout(object):
 
