feature/wrappers-modify

1 files changed, 32 insertions, 54 deletions

diff --git a/src/mailman_pgp/testing/pgp.py b/src/mailman_pgp/testing/pgp.py
index afb62c8..a50ae02 100644
--- a/src/mailman_pgp/testing/pgp.py
+++ b/src/mailman_pgp/testing/pgp.py
@@ -88,16 +88,12 @@ class WrapperTestCase(unittest.TestCase):
         self.assertEqual(wrapped.has_signature(), has)
 
     def sign(self, message, key):
-        wrapped = self.wrap(message)
-        signed = wrapped.sign(key)
-        signed_wrapped = self.wrap(signed)
-        self.assertTrue(signed_wrapped.is_signed())
+        signed = self.wrap(message).sign(key)
+        self.assertTrue(signed.is_signed())
 
     def sign_verify(self, message, priv, pub):
-        wrapped = self.wrap(message)
-        signed = wrapped.sign(priv)
-        signed_wrapped = self.wrap(signed)
-        for signature in signed_wrapped.verify(pub):
+        signed = self.wrap(message).sign(priv)
+        for signature in signed.verify(pub):
             self.assertTrue(bool(signature))
 
     def verify(self, message, key, valid):
@@ -114,29 +110,20 @@ class WrapperTestCase(unittest.TestCase):
         self.assertEqual(wrapped.has_encryption(), has)
 
     def encrypt(self, message, *keys, **kwargs):
-        wrapped = self.wrap(message)
-        encrypted = wrapped.encrypt(*keys, **kwargs)
-        encrypted_wrapped = self.wrap(encrypted)
-        self.assertTrue(encrypted_wrapped.is_encrypted())
+        encrypted = self.wrap(message).encrypt(*keys, **kwargs)
+        self.assertTrue(encrypted.is_encrypted())
 
     def encrypt_decrypt(self, message, pub, priv):
-        wrapped = self.wrap(message)
-        encrypted = wrapped.encrypt(pub)
+        decrypted = self.wrap(message).copy().encrypt(pub).decrypt(priv)
 
-        encrypted_wrapped = self.wrap(encrypted)
-        decrypted = encrypted_wrapped.decrypt(priv)
-        decrypted_wrapped = self.wrap(decrypted)
-
-        self.assertFalse(decrypted_wrapped.is_encrypted())
-        self.assertTrue(payload_equal(decrypted, message))
+        self.assertFalse(decrypted.is_encrypted())
+        self.assertTrue(payload_equal(decrypted.msg, message))
 
     def decrypt(self, message, key, clear):
-        wrapped = self.wrap(message)
-        decrypted = wrapped.decrypt(key)
-        decrypted_wrapped = self.wrap(decrypted)
+        decrypted = self.wrap(message).copy().decrypt(key)
 
-        self.assertFalse(decrypted_wrapped.is_encrypted())
-        self.assertEqual(decrypted.get_payload(), clear)
+        self.assertFalse(decrypted.is_encrypted())
+        self.assertEqual(decrypted.msg.get_payload(), clear)
 
     def has_keys(self, message, has_keys):
         wrapped = self.wrap(message)
@@ -156,9 +143,7 @@ class WrapperTestCase(unittest.TestCase):
         self.assertListEqual(loaded_fingerprints, fingerprints)
 
     def attach_keys(self, message, keys):
-        wrapped = self.wrap(message)
-        attached = wrapped.attach_keys(*keys)
-        wrapped = self.wrap(attached)
+        wrapped = self.wrap(message).attach_keys(*keys)
         loaded = list(wrapped.keys())
 
         self.assertTrue(wrapped.has_keys())
@@ -184,9 +169,7 @@ class WrapperTestCase(unittest.TestCase):
         self.assertListEqual(loaded_issuers, issuers)
 
     def attach_revocs(self, message, revocs):
-        wrapped = self.wrap(message)
-        attached = wrapped.attach_revocs(*revocs)
-        wrapped = self.wrap(attached)
+        wrapped = self.wrap(message).attach_revocs(*revocs)
         loaded = list(wrapped.revocs())
 
         self.assertTrue(wrapped.has_revocs())
@@ -196,46 +179,41 @@ class WrapperTestCase(unittest.TestCase):
 
     def sign_encrypt_decrypt_verify(self, message, sign_key, encrypt_key):
         wrapped = self.wrap(message)
-        encrypted = wrapped.sign_encrypt(sign_key, encrypt_key.pubkey)
-        encrypted_wrapped = self.wrap(encrypted)
-        self.assertTrue(encrypted_wrapped.is_encrypted())
+        encrypted = wrapped.copy().sign_encrypt(sign_key, encrypt_key.pubkey)
+        self.assertTrue(encrypted.is_encrypted())
 
-        decrypted = encrypted_wrapped.decrypt(encrypt_key)
-        decrypted_wrapped = self.wrap(decrypted)
-        self.assertTrue(decrypted_wrapped.is_signed())
-        self.assertFalse(decrypted_wrapped.is_encrypted())
+        decrypted = encrypted.decrypt(encrypt_key)
+        self.assertTrue(decrypted.is_signed())
+        self.assertFalse(decrypted.is_encrypted())
 
-        verification = decrypted_wrapped.verify(sign_key.pubkey)
+        verification = decrypted.verify(sign_key.pubkey)
         for sig in verification:
             self.assertTrue(bool(sig))
-        self.assertListEqual(list(decrypted_wrapped.get_signed()),
+        self.assertListEqual(list(decrypted.get_signed()),
                              list(wrapped.get_payload()))
 
     def sign_then_encrypt_decrypt_verify(self, message, sign_key, encrypt_key):
         wrapped = self.wrap(message)
-        encrypted = wrapped.sign_then_encrypt(sign_key, encrypt_key.pubkey)
-        encrypted_wrapped = self.wrap(encrypted)
-        self.assertTrue(encrypted_wrapped.is_encrypted())
+        encrypted = wrapped.copy().sign(sign_key).encrypt(encrypt_key.pubkey)
+        self.assertTrue(encrypted.is_encrypted())
 
-        decrypted = encrypted_wrapped.decrypt(encrypt_key)
-        decrypted_wrapped = self.wrap(decrypted)
-        self.assertTrue(decrypted_wrapped.is_signed())
-        self.assertFalse(decrypted_wrapped.is_encrypted())
+        decrypted = encrypted.copy().decrypt(encrypt_key)
+        self.assertTrue(decrypted.is_signed())
+        self.assertFalse(decrypted.is_encrypted())
 
-        verification = decrypted_wrapped.verify(sign_key.pubkey)
+        verification = decrypted.verify(sign_key.pubkey)
         for sig in verification:
             self.assertTrue(bool(sig))
-        self.assertListEqual(list(decrypted_wrapped.get_signed()),
+        self.assertListEqual(list(decrypted.get_signed()),
                              list(wrapped.get_payload()))
 
     def decrypt_verify(self, message, decrypt_key, verify_key, valid):
         wrapped = self.wrap(message)
-        decrypted = wrapped.decrypt(decrypt_key)
-        decrypted_wrapped = self.wrap(decrypted)
+        decrypted = wrapped.copy().decrypt(decrypt_key)
 
-        self.assertFalse(decrypted_wrapped.is_encrypted())
-        self.assertTrue(decrypted_wrapped.is_signed())
+        self.assertFalse(decrypted.is_encrypted())
+        self.assertTrue(decrypted.is_signed())
 
-        verification = decrypted_wrapped.verify(verify_key)
+        verification = decrypted.verify(verify_key)
         for sig in verification:
             self.assertEqual(bool(sig), valid)