From be556d24d4771725d67905e00cdb1526119a21ba Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Viktoras=20Laukevi=C4=8Dius?=
 <viktoras.laukevicius@yahoo.com>
Date: Wed, 14 Aug 2019 16:16:02 +0300
Subject: [PATCH] Add support for arbitrary key-value pairs in
 messaging.ApsAlert

---
 firebase_admin/_messaging_utils.py | 14 ++++++-
 tests/test_messaging.py            | 66 ++++++++++++++++++++++++++++++
 2 files changed, 79 insertions(+), 1 deletion(-)

diff --git a/firebase_admin/_messaging_utils.py b/firebase_admin/_messaging_utils.py
index 09f7daf87..d6e263dcf 100644
--- a/firebase_admin/_messaging_utils.py
+++ b/firebase_admin/_messaging_utils.py
@@ -394,10 +394,13 @@ class ApsAlert(object):
         action_loc_key: Key of the text in the app's string resources to use to localize the
             action button text (optional).
         launch_image: Image for the notification action (optional).
+        custom_data: A dict of custom key-value pairs to be included in the ApsAlert dictionary
+            (optional)
     """
 
     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 +410,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 +839,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
index 878e1365b..67ee0721d 100644
--- 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):