refactor tests, pytest style

3 files changed, 79 insertions, 86 deletions

diff --git a/requirements.pip b/requirements.pip
index d1a7d18..b3d5230 100644
--- a/requirements.pip
+++ b/requirements.pip
@@ -1,6 +1,3 @@
-six==1.10.0 \
-    --hash=sha256:0ff78c403d9bccf5a425a6d31a12aa6b47f1c21ca4dc2573a7e2f32a97335eb1 \
-    --hash=sha256:105f8d68616f8248e24bf0e9372ef04d3cc10104f1980f54d57b2ce73a5ad56a
 gnupg==2.0.2 \
     --hash=sha256:2a7f287b4bf2613513be6fe66810f6c43a89092af2e99406986eba12883ec602 \
     --hash=sha256:17acbace8a51d86014748e0dab7fcef962cea33451934528e401c546152be9ec \
diff --git a/tests/test_protection.py b/tests/test_protection.py
index e84b15b..70c1a57 100644
--- a/tests/test_protection.py
+++ b/tests/test_protection.py
@@ -1,5 +1,4 @@
 import six
-import unittest
 from base64 import b64encode
 from email.parser import Parser
 from zope.interface import implementer
@@ -23,85 +22,85 @@ Subject: %(subject)s
     "body": BODY
 }
 
+parser = Parser()
 
-class ProtectTest(unittest.TestCase):
-    def test_pgp_encrypted_mime(self):
-        p = Parser()
-        msg = p.parsestr(EMAIL)
-        encrypter = Encrypter()
-        conf = ProtectConfig(openpgp=encrypter, obscured_headers=[])
-        encmsg = protect(msg, config=conf)
-
-        self.assertEqual(encmsg.get_payload(1).get_payload(), encrypter.encstr)
-        self._assert_body(encrypter.data, BODY+'\n')
-        self.assertEqual([TO], encrypter.encraddr)
-        self.assertEqual(encmsg.get_content_type(), "multipart/encrypted")
-
-    def test_unobscured_headers(self):
-        p = Parser()
-        msg = p.parsestr(EMAIL)
-        encrypter = Encrypter()
-        conf = ProtectConfig(openpgp=encrypter, obscured_headers=[])
-        encmsg = protect(msg, config=conf)
-
-        self.assertEqual(encmsg['from'], FROM)
-        self.assertEqual(encmsg['to'], TO)
-        self.assertEqual(encmsg['subject'], SUBJECT)
-
-    def test_obscured_headers(self):
-        p = Parser()
-        msg = p.parsestr(EMAIL)
-        encrypter = Encrypter()
-        conf = ProtectConfig(openpgp=encrypter)
-        encmsg = protect(msg, config=conf)
-
-        for header, value in conf.obscured_headers.items():
-            msgheaders = encmsg.get_all(header, [])
-            if msgheaders:
-                self.assertEqual(msgheaders, [value])
-
-        encpart = p.parsestr(encrypter.data)
-        self.assertEqual(encpart.get_content_type(), "multipart/mixed")
-        rfc822part = encpart.get_payload(0)
-        self.assertEqual(rfc822part.get_content_type(), "text/rfc822-headers")
-        rfc822body = "Subject: %s\n" % (SUBJECT,)
-        self.assertEqual(rfc822part.get_payload(), rfc822body)
-        self.assertEqual(encpart.get_payload(1).get_payload(),
-                         BODY+'\n')
-
-    def test_pgp_signed_mime(self):
-        p = Parser()
-        msg = p.parsestr(EMAIL)
-        signer = Signer()
-        conf = ProtectConfig(openpgp=signer)
-        encmsg = protect(msg, encrypt=False, config=conf)
-
-        b64body = b64encode(six.b(BODY+'\n'))
-        self.assertEqual(six.b(encmsg.get_payload(0).get_payload()), b64body)
-        self.assertEqual(encmsg.get_payload(1).get_payload(), signer.signature)
-        self._assert_body(signer.data, b64body.decode('utf-8'))
-        self.assertEqual(encmsg.get_content_type(), "multipart/signed")
-
-    def test_signed_headers(self):
-        p = Parser()
-        msg = p.parsestr(EMAIL)
-        signer = Signer()
-        conf = ProtectConfig(openpgp=signer)
-        signmsg = protect(msg, encrypt=False, config=conf)
-
-        self.assertEqual(signmsg['from'], FROM)
-        self.assertEqual(signmsg['to'], TO)
-        self.assertEqual(signmsg['subject'], SUBJECT)
-
-        signedpart = signmsg.get_payload(0)
-        self.assertEqual(signedpart['from'], FROM)
-        self.assertEqual(signedpart['to'], TO)
-        self.assertEqual(signedpart['subject'], SUBJECT)
-
-    def _assert_body(self, data, body):
-        p = Parser()
-        msg = p.parsestr(data)
-        self.assertEqual(msg.get_payload(), body)
+
+def test_pgp_encrypted_mime():
+    msg = parser.parsestr(EMAIL)
+    encrypter = Encrypter()
+    conf = ProtectConfig(openpgp=encrypter, obscured_headers=[])
+    encmsg = protect(msg, config=conf)
+
+    assert encmsg.get_payload(1).get_payload() == encrypter.encstr
+    assert [TO] == encrypter.encraddr
+    assert encmsg.get_content_type() == "multipart/encrypted"
+    assert get_body(encrypter.data) == BODY + '\n'
+
+
+def test_unobscured_headers():
+    msg = parser.parsestr(EMAIL)
+    encrypter = Encrypter()
+    conf = ProtectConfig(openpgp=encrypter, obscured_headers=[])
+    encmsg = protect(msg, config=conf)
+
+    assert encmsg['from'] == FROM
+    assert encmsg['to'] == TO
+    assert encmsg['subject'] == SUBJECT
+
+
+def test_obscured_headers():
+    msg = parser.parsestr(EMAIL)
+    encrypter = Encrypter()
+    conf = ProtectConfig(openpgp=encrypter)
+    encmsg = protect(msg, config=conf)
+
+    for header, value in conf.obscured_headers.items():
+        msgheaders = encmsg.get_all(header, [])
+        if msgheaders:
+            assert msgheaders == [value]
+
+    encpart = parser.parsestr(encrypter.data)
+    assert encpart.get_content_type() == "multipart/mixed"
+
+    rfc822part = encpart.get_payload(0)
+    assert rfc822part.get_content_type() == "text/rfc822-headers"
+
+    rfc822body = "Subject: %s\n" % (SUBJECT,)
+    assert rfc822part.get_payload() == rfc822body
+    assert encpart.get_payload(1).get_payload() == BODY+'\n'
+
+
+def test_pgp_signed_mime():
+    msg = parser.parsestr(EMAIL)
+    signer = Signer()
+    conf = ProtectConfig(openpgp=signer)
+    encmsg = protect(msg, encrypt=False, config=conf)
+
+    b64body = b64encode(six.b(BODY+'\n'))
+    assert six.b(encmsg.get_payload(0).get_payload()) == b64body
+    assert encmsg.get_payload(1).get_payload() == signer.signature
+    assert get_body(signer.data) == b64body.decode('utf-8')
+    assert encmsg.get_content_type() == "multipart/signed"
+
+
+def test_signed_headers():
+    msg = parser.parsestr(EMAIL)
+    signer = Signer()
+    conf = ProtectConfig(openpgp=signer)
+    signmsg = protect(msg, encrypt=False, config=conf)
+
+    assert signmsg['from'] == FROM
+    assert signmsg['to'] == TO
+    assert signmsg['subject'] == SUBJECT
+
+    signedpart = signmsg.get_payload(0)
+    assert signedpart['from'] == FROM
+    assert signedpart['to'] == TO
+    assert signedpart['subject'] == SUBJECT
+
+
+def get_body(data):
+    return parser.parsestr(data).get_payload()
 
 
 @implementer(IOpenPGP)
@@ -121,7 +120,3 @@ class Signer(object):
     def sign(self, data):
         self.data = data
         return self.signature
-
-
-if __name__ == "__main__":
-    unittest.main()
diff --git a/tox.ini b/tox.ini
index bc5335d..5dd88f0 100644
--- a/tox.ini
+++ b/tox.ini
@@ -13,6 +13,7 @@ deps =
     pytest
     pytest-pep8
     pytest-flake8
+    pytest-sugar
     six
     gnupg
     zope.interface