refactor benchmark tests

Author: purificant

tests/test_benchmark.py

@@ -1,8 +1,10 @@
 """ This module contains benchmark tests intended to guide development of a performant codebase. """
+import nacl
+import pysodium
 import pytest
 
-from paseto.protocol.version2 import Version2
 from paseto.crypto import primitives
+from paseto.protocol.version2 import Version2
 
 KEY = b"0" * 32
 MESSAGE = b"foo"
@@ -12,18 +14,17 @@
 @pytest.mark.benchmark(group="encrypt")
 def test_encrypt_one(benchmark):
     """Benchmark only encryption."""
-    from nacl.bindings import crypto_aead_xchacha20poly1305_ietf_encrypt as encrypt
-    primitives.encrypt = encrypt
+    primitives.encrypt = nacl.bindings.crypto_aead_xchacha20poly1305_ietf_encrypt
 
     token = benchmark(Version2.encrypt, MESSAGE, KEY, FOOTER)
     plain_text = Version2.decrypt(token, KEY, FOOTER)
     assert plain_text == MESSAGE
 
+
 @pytest.mark.benchmark(group="encrypt")
 def test_encrypt_two(benchmark):
     """Benchmark only encryption."""
-    from pysodium import crypto_aead_xchacha20poly1305_ietf_encrypt as encrypt
-    primitives.encrypt = encrypt
+    primitives.encrypt = pysodium.crypto_aead_xchacha20poly1305_ietf_encrypt
 
     token = benchmark(Version2.encrypt, MESSAGE, KEY, FOOTER)
     plain_text = Version2.decrypt(token, KEY, FOOTER)
@@ -33,8 +34,7 @@ def test_encrypt_two(benchmark):
 @pytest.mark.benchmark(group="decrypt")
 def test_decrypt_one(benchmark):
     """Benchmark only decryption."""
-    from nacl.bindings import crypto_aead_xchacha20poly1305_ietf_decrypt as decrypt
-    primitives.decrypt = decrypt
+    primitives.decrypt = nacl.bindings.crypto_aead_xchacha20poly1305_ietf_decrypt
 
     token = Version2.encrypt(MESSAGE, KEY, FOOTER)
     plain_text = benchmark(Version2.decrypt, token, KEY, FOOTER)
@@ -44,20 +44,18 @@ def test_decrypt_one(benchmark):
 @pytest.mark.benchmark(group="decrypt")
 def test_decrypt_two(benchmark):
     """Benchmark only decryption."""
-    from pysodium import crypto_aead_xchacha20poly1305_ietf_decrypt as decrypt
-    primitives.decrypt = decrypt
+    primitives.decrypt = pysodium.crypto_aead_xchacha20poly1305_ietf_decrypt
 
     token = Version2.encrypt(MESSAGE, KEY, FOOTER)
     plain_text = benchmark(Version2.decrypt, token, KEY, FOOTER)
     assert plain_text == MESSAGE
 
+
 @pytest.mark.benchmark(group="encrypt_and_decrypt")
 def test_encrypt_and_decrypt_one(benchmark):
     """Benchmark encryption and decryption run together."""
-    from nacl.bindings import crypto_aead_xchacha20poly1305_ietf_encrypt as encrypt
-    primitives.encrypt = encrypt
-    from nacl.bindings import crypto_aead_xchacha20poly1305_ietf_decrypt as decrypt
-    primitives.decrypt = decrypt
+    primitives.encrypt = nacl.bindings.crypto_aead_xchacha20poly1305_ietf_encrypt
+    primitives.decrypt = nacl.bindings.crypto_aead_xchacha20poly1305_ietf_decrypt
 
     def encrypt_and_decrypt():
         token = Version2.encrypt(MESSAGE, KEY, FOOTER)
@@ -66,13 +64,12 @@ def encrypt_and_decrypt():
     plain_text = benchmark(encrypt_and_decrypt)
     assert plain_text == MESSAGE
 
+
 @pytest.mark.benchmark(group="encrypt_and_decrypt")
 def test_encrypt_and_decrypt_two(benchmark):
     """Benchmark encryption and decryption run together."""
-    from pysodium import crypto_aead_xchacha20poly1305_ietf_decrypt as decrypt
-    primitives.decrypt = decrypt
-    from pysodium import crypto_aead_xchacha20poly1305_ietf_encrypt as encrypt
-    primitives.encrypt = encrypt
+    primitives.decrypt = pysodium.crypto_aead_xchacha20poly1305_ietf_decrypt
+    primitives.encrypt = pysodium.crypto_aead_xchacha20poly1305_ietf_encrypt
 
     def encrypt_and_decrypt():
         token = Version2.encrypt(MESSAGE, KEY, FOOTER)