tests: reduce generate_key usage

Author: lepture

tests/jwe/test_json.py

@@ -10,12 +10,13 @@
 
 
 class TestJWEJSON(TestCase):
+    rsa_key = RSAKey.generate_key()
+    ec_key = ECKey.generate_key()
+
     def test_multiple_recipients_with_key(self):
-        key1 = RSAKey.generate_key()
-        key2 = ECKey.generate_key()
         obj = GeneralJSONEncryption({"enc": "A128CBC-HS256"}, b"i")
-        obj.add_recipient({"alg": "RSA-OAEP"}, key1)
-        obj.add_recipient({"alg": "ECDH-ES+A128KW"}, key2)
+        obj.add_recipient({"alg": "RSA-OAEP"}, self.rsa_key)
+        obj.add_recipient({"alg": "ECDH-ES+A128KW"}, self.ec_key)
         value = jwe.encrypt_json(obj, None)
         self.assertIn("recipients", value)
         self.assertEqual(len(value["recipients"]), 2)
@@ -33,7 +34,7 @@ def test_multiple_recipients_without_key(self):
     def test_multiple_recipients_with_direct_mode(self):
         obj = GeneralJSONEncryption({"enc": "A128CBC-HS256"}, b"i")
         obj.add_recipient({"alg": "dir"}, OctKey.generate_key())
-        obj.add_recipient({"alg": "RSA-OAEP"}, RSAKey.generate_key())
+        obj.add_recipient({"alg": "RSA-OAEP"}, self.rsa_key)
         self.assertRaises(
             ConflictAlgorithmError,
             jwe.encrypt_json,
@@ -42,26 +43,23 @@ def test_multiple_recipients_with_direct_mode(self):
 
     def test_with_aad(self):
         obj = GeneralJSONEncryption({"enc": "A128CBC-HS256"}, b"i", aad=b"foo")
-        key1 = RSAKey.generate_key()
-        obj.add_recipient({"alg": "RSA-OAEP"}, key1)
+        obj.add_recipient({"alg": "RSA-OAEP"}, self.rsa_key)
         value = jwe.encrypt_json(obj, None)
-        obj1 = jwe.decrypt_json(value, key1)
+        obj1 = jwe.decrypt_json(value, self.rsa_key)
         self.assertEqual(obj1.aad, b"foo")
 
     def test_decode_multiple_recipients(self):
-        key1 = RSAKey.generate_key()
-        key2 = ECKey.generate_key()
         obj = GeneralJSONEncryption({"enc": "A128CBC-HS256"}, b"i")
-        obj.add_recipient({"alg": "RSA-OAEP"}, key1)
-        obj.add_recipient({"alg": "ECDH-ES+A128KW"}, key2)
+        obj.add_recipient({"alg": "RSA-OAEP"}, self.rsa_key)
+        obj.add_recipient({"alg": "ECDH-ES+A128KW"}, self.ec_key)
         value = jwe.encrypt_json(obj, None)
         self.assertRaises(
             InvalidKeyTypeError,
             jwe.decrypt_json,
-            value, key1,
+            value, self.rsa_key,
         )
         registry = jwe.JWERegistry(verify_all_recipients=False)
-        obj1 = jwe.decrypt_json(value, key1, registry=registry)
+        obj1 = jwe.decrypt_json(value, self.rsa_key, registry=registry)
         self.assertEqual(obj1.plaintext, b"i")
 
         key3 = OctKey.generate_key()

tests/jwk/test_ec_key.py

@@ -5,6 +5,8 @@
 
 
 class TestECKey(TestCase):
+    default_key = ECKey.generate_key()
+
     def test_exchange_derive_key(self):
         key1 = ECKey.generate_key("P-256")
         key2 = ECKey.generate_key("P-384")
@@ -58,9 +60,8 @@ def test_generate_key(self):
         self.assertIsNotNone(key.kid)
 
     def test_import_from_der_bytes(self):
-        origin_key = ECKey.generate_key()
-        value1 = origin_key.as_der()
-        value2 = origin_key.as_der(private=False)
+        value1 = self.default_key.as_der()
+        value2 = self.default_key.as_der(private=False)
 
         key1 = ECKey.import_key(value1)
         key2 = ECKey.import_key(value2)
@@ -80,13 +81,11 @@ def test_output_with_password(self):
         self.assertEqual(key.as_dict(), key2.as_dict())
 
     def test_key_eq(self):
-        key1 = ECKey.generate_key()
+        key1 = self.default_key
         key2 = ECKey.import_key(key1.as_dict())
         self.assertEqual(key1, key2)
         key3 = ECKey.generate_key()
         self.assertNotEqual(key1, key3)
 
     def test_key_eq_with_different_types(self):
-        key1 = ECKey.generate_key()
-        key2 = OctKey.generate_key()
-        self.assertNotEqual(key1, key2)
+        self.assertNotEqual(self.default_key, OctKey.generate_key())

tests/jwk/test_rsa_key.py

@@ -4,6 +4,8 @@
 
 
 class TestRSAKey(TestCase):
+    default_key = RSAKey.generate_key()
+
     def test_import_key_from_dict(self):
         # https://www.rfc-editor.org/rfc/rfc7517#appendix-A.1
         data = {
@@ -128,8 +130,7 @@ def test_generate_key(self):
         self.assertIsNotNone(key.kid)
 
     def test_import_from_der_bytes(self):
-        origin_key = RSAKey.generate_key()
-        value1 = origin_key.as_der()
+        value1 = self.default_key.as_der()
         key1 = RSAKey.import_key(value1)
         self.assertEqual(value1, key1.as_der())
 
@@ -152,7 +153,7 @@ def test_output_with_password(self):
         self.assertEqual(key.as_dict(), key2.as_dict())
 
     def test_key_eq(self):
-        key1 = RSAKey.generate_key()
+        key1 = self.default_key
         key2 = RSAKey.import_key(key1.as_dict())
         self.assertIsNot(key1, key2)
         self.assertEqual(key1, key2)