Merge pull request #1 from canwehatch/extend_PayloadAlert

Extend PayloadAlert with all keys supported

Author: moden-py

apns.py

@@ -268,36 +268,61 @@ def write(self, string):
 
 
 class PayloadAlert(object):
-    def __init__(self, body=None, title = None, subtitle = None, action_loc_key=None, loc_key=None,
-                 loc_args=None, launch_image=None):
-        super(PayloadAlert, self).__init__()
-        
+    """
+    Payload for APNS alert.
+    https://developer.apple.com/library/content/documentation/NetworkingInternet/Conceptual/RemoteNotificationsPG/PayloadKeyReference.html
+    """
+    def __init__(self,
+                 body=None,
+                 title=None,
+                 subtitle=None,
+                 action_loc_key=None,
+                 loc_key=None,
+                 loc_args=None,
+                 launch_image=None,
+                 title_loc_key=None,
+                 title_loc_args=None):
+
         self.body = body
         self.title = title
         self.subtitle = subtitle
         self.action_loc_key = action_loc_key
         self.loc_key = loc_key
         self.loc_args = loc_args
         self.launch_image = launch_image
+        self.title_loc_key = title_loc_key
+        self.title_loc_args = title_loc_args
+
+        self._dict = {
+            'body': self.body,
+            'title': self.title,
+            'subtitle': self.subtitle,
+            'action-loc-key': self.action_loc_key,
+            'loc-key': self.loc_key,
+            'loc-args': self.loc_args,
+            'launch-image': self.launch_image,
+            'title-loc-key': self.title_loc_key,
+            'title-loc-args': self.title_loc_args
+        }
 
     def dict(self):
-        d = {}
-        
-        if self.body:
-            d['body'] = self.body
-        if self.title:
-            d['title'] = self.title
-        if self.subtitle:
-            d['subtitle'] = self.subtitle
-        if self.action_loc_key:
-            d['action-loc-key'] = self.action_loc_key
-        if self.loc_key:
-            d['loc-key'] = self.loc_key
-        if self.loc_args:
-            d['loc-args'] = self.loc_args
-        if self.launch_image:
-            d['launch-image'] = self.launch_image
-        return d
+        cleared = {
+            key: value
+            for (key, value)
+            in self._dict.items()
+            if value is not None
+        }
+        return cleared
+
+    def __eq__(self, other):
+        return self.dict() == other.dict()
+
+    def __ne__(self, other):
+        return self.dict() != other.dict()
+
+    def __repr__(self):
+        return 'PayloadAlert(**{})'.format(self.dict())
+
 
 class PayloadTooLargeError(Exception):
     def __init__(self, payload_size):

tests.py

@@ -188,9 +188,19 @@ def testFrame(self):
         frame = Frame()
         frame.add_item(token_hex, payload, identifier, expiry, priority)
 
-        f1 = bytearray(b'\x02\x00\x00\x00t\x01\x00 \xb5\xbb\x9d\x80\x14\xa0\xf9\xb1\xd6\x1e!\xe7\x96\xd7\x8d\xcc\xdf\x13R\xf2<\xd3(\x12\xf4\x85\x0b\x87\x8a\xe4\x94L\x02\x00<{"aps":{"sound":"default","badge":4,"alert":"Hello World!"}}\x03\x00\x04\x00\x00\x00\x01\x04\x00\x04\x00\x00\x0e\x10\x05\x00\x01\n')
-        f2 = bytearray(b'\x02\x00\x00\x00t\x01\x00 \xb5\xbb\x9d\x80\x14\xa0\xf9\xb1\xd6\x1e!\xe7\x96\xd7\x8d\xcc\xdf\x13R\xf2<\xd3(\x12\xf4\x85\x0b\x87\x8a\xe4\x94L\x02\x00<{"aps":{"sound":"default","alert":"Hello World!","badge":4}}\x03\x00\x04\x00\x00\x00\x01\x04\x00\x04\x00\x00\x0e\x10\x05\x00\x01\n')
-        self.assertTrue(f1 == frame.get_frame() or f2 == frame.get_frame())
+        frame_bytes = frame.get_frame()
+
+        prefix = frame_bytes[:43]
+        data = frame_bytes[43: -18]
+        postfix = frame_bytes[-18:]
+
+        self.assertEqual(prefix, b'\x02\x00\x00\x00t\x01\x00 \xb5\xbb\x9d\x80\x14\xa0\xf9\xb1\xd6\x1e!\xe7\x96\xd7\x8d\xcc\xdf\x13R\xf2<\xd3(\x12\xf4\x85\x0b\x87\x8a\xe4\x94L\x02\x00<')
+        self.assertEqual(
+            json.loads(data.decode()),
+            json.loads('{"aps":{"sound":"default","badge":4,'
+                       '"alert":"Hello World!"}}')
+        )
+        self.assertEqual(postfix, b'\x03\x00\x04\x00\x00\x00\x01\x04\x00\x04\x00\x00\x0e\x10\x05\x00\x01\n')
 
     def testPayloadTooLargeError(self):
         # The maximum size of the JSON payload is MAX_PAYLOAD_LENGTH