Add tests for get_signature and get_encrypted.

1 files changed, 24 insertions, 1 deletions

diff --git a/src/mailman_pgp/pgp/tests/base.py b/src/mailman_pgp/pgp/tests/base.py
index 47d333c..e16c95e 100644
--- a/src/mailman_pgp/pgp/tests/base.py
+++ b/src/mailman_pgp/pgp/tests/base.py
@@ -41,9 +41,24 @@ def load_key(path):
     return key
 
 
+def payload_equal(one_msg, other_msg):
+    one_payload = one_msg.get_payload()
+    other_payload = other_msg.get_payload()
+    if isinstance(one_payload, list) and isinstance(other_payload, list):
+        if len(one_payload) != len(other_payload):
+            return False
+        for one_inner, other_inner in zip(one_payload, other_payload):
+            if not payload_equal(one_inner, other_inner):
+                return False
+        return True
+    else:
+        return one_payload == other_payload
+
+
 class WrapperTestCase(TestCase):
     layer = PGPLayer
     wrapper = None
+    maxDiff = None
 
     def wrap(self, message):
         return self.wrapper(message)
@@ -52,6 +67,10 @@ class WrapperTestCase(TestCase):
         wrapped = self.wrap(message)
         self.assertEqual(wrapped.is_signed(), signed)
 
+    def has_signature(self, message, has):
+        wrapped = self.wrap(message)
+        self.assertEqual(wrapped.has_signature(), has)
+
     def sign(self, message, key):
         wrapped = self.wrap(message)
         signed = wrapped.sign(key)
@@ -74,6 +93,10 @@ class WrapperTestCase(TestCase):
         wrapped = self.wrap(message)
         self.assertEqual(wrapped.is_encrypted(), encrypted)
 
+    def has_encryption(self, message, has):
+        wrapped = self.wrap(message)
+        self.assertEqual(wrapped.has_encryption(), has)
+
     def encrypt(self, message, *keys, **kwargs):
         wrapped = self.wrap(message)
         encrypted = wrapped.encrypt(*keys, **kwargs)
@@ -89,7 +112,7 @@ class WrapperTestCase(TestCase):
         decrypted_wrapped = self.wrap(decrypted)
 
         self.assertFalse(decrypted_wrapped.is_encrypted())
-        self.assertEqual(decrypted.get_payload(), message.get_payload())
+        self.assertTrue(payload_equal(decrypted, message))
 
     def decrypt(self, message, key, clear):
         wrapped = self.wrap(message)