1 files changed, 63 insertions, 8 deletions

diff --git a/src/mailman_pgp/testing/pgp.py b/src/mailman_pgp/testing/pgp.py
index 844b708..89f7a98 100644
--- a/src/mailman_pgp/testing/pgp.py
+++ b/src/mailman_pgp/testing/pgp.py
@@ -19,6 +19,7 @@
 
 import os
 from email import message_from_bytes
+from operator import attrgetter
 from unittest import TestCase
 
 from mailman.email.message import Message
@@ -26,21 +27,36 @@ from pgpy import PGPKey
 from pkg_resources import resource_string
 
 from mailman_pgp.testing.layers import PGPLayer
+from mailman_pgp.utils.pgp import revoc_from_blob
+
+
+def load_blob(*path):
+    return resource_string('mailman_pgp.pgp.tests',
+                           os.path.join('data', *path))
 
 
 def load_message(path):
-    data = resource_string('mailman_pgp.pgp.tests',
-                           os.path.join('data', 'messages', path))
-    return message_from_bytes(data, Message)
+    """
+    :rtype: Message
+    """
+    return message_from_bytes(load_blob('messages', path), Message)
 
 
 def load_key(path):
-    key, _ = PGPKey.from_blob(
-            resource_string('mailman_pgp.pgp.tests',
-                            os.path.join('data', 'keys', path)))
+    """
+    :rtype: pgpy.PGPKey
+    """
+    key, _ = PGPKey.from_blob(load_blob('keys', path))
     return key
 
 
+def load_revoc(path):
+    """
+    :rtype: pgpy.PGPSignature
+    """
+    return revoc_from_blob(load_blob('revocs', path))
+
+
 def payload_equal(one_msg, other_msg):
     one_payload = one_msg.get_payload()
     other_payload = other_msg.get_payload()
@@ -135,10 +151,49 @@ class WrapperTestCase(TestCase):
         loaded = list(wrapped.keys())
         self.assertEqual(len(loaded), len(keys))
 
-        loaded_fingerprints = list(map(lambda key: key.fingerprint, loaded))
-        fingerprints = list(map(lambda key: key.fingerprint, keys))
+        loaded_fingerprints = list(map(attrgetter('fingerprint'), loaded))
+        fingerprints = list(map(attrgetter('fingerprint'), keys))
+        self.assertListEqual(loaded_fingerprints, fingerprints)
+
+    def attach_keys(self, message, keys):
+        wrapped = self.wrap(message)
+        attached = wrapped.attach_keys(*keys)
+        wrapped = self.wrap(attached)
+        loaded = list(wrapped.keys())
+
+        self.assertTrue(wrapped.has_keys())
+        loaded_fingerprints = list(map(attrgetter('fingerprint'), loaded))
+        fingerprints = list(map(attrgetter('fingerprint'), keys))
         self.assertListEqual(loaded_fingerprints, fingerprints)
 
+    def has_revocs(self, message, has_revocs):
+        wrapped = self.wrap(message)
+        self.assertEqual(wrapped.has_revocs(), has_revocs)
+
+    def is_revocs(self, message, is_revocs):
+        wrapped = self.wrap(message)
+        self.assertEqual(wrapped.is_revocs(), is_revocs)
+
+    def revocs(self, message, revocs):
+        wrapped = self.wrap(message)
+        loaded = list(wrapped.revocs())
+        self.assertEqual(len(loaded), len(revocs))
+
+        loaded_issuers = list(map(attrgetter('signer'), loaded))
+        issuers = list(map(attrgetter('signer'), revocs))
+        self.assertListEqual(loaded_issuers, issuers)
+
+    def attach_revocs(self, message, revocs):
+        wrapped = self.wrap(message)
+        attached = wrapped.attach_revocs(*revocs)
+        wrapped = self.wrap(attached)
+        loaded = list(wrapped.revocs())
+
+        self.assertTrue(wrapped.has_revocs())
+        loaded_issuers = list(map(attrgetter('signer'), loaded))
+        issuers = list(map(attrgetter('signer'), revocs))
+        self.assertListEqual(loaded_issuers, issuers)
+
     def sign_encrypt_decrypt_verify(self, message, sign_key, encrypt_key):
         wrapped = self.wrap(message)
         encrypted = wrapped.sign_encrypt(sign_key, encrypt_key.pubkey)