1 files changed, 80 insertions, 17 deletions
diff --git a/tests/hazmat/primitives/twofactor/test_totp.py b/tests/hazmat/primitives/twofactor/test_totp.py
index 58819be4..8877a70a 100644
--- a/tests/hazmat/primitives/twofactor/test_totp.py
+++ b/tests/hazmat/primitives/twofactor/test_totp.py
@@ -15,7 +15,6 @@ import pytest
 
 from cryptography.exceptions import InvalidToken
 from cryptography.hazmat.primitives import hashes
-from cryptography.hazmat.primitives.hashes import SHA1
 from cryptography.hazmat.primitives.twofactor.totp import TOTP
 from tests.utils import load_vectors_from_file, load_nist_vectors
 
@@ -26,42 +25,106 @@ vectors = load_vectors_from_file(
 @pytest.mark.hmac
 class TestTOTP(object):
 
-    @pytest.mark.parametrize("params", vectors)
-    def test_generate(self, backend, params):
+    @pytest.mark.supported(
+        only_if=lambda backend: backend.hmac_supported(hashes.SHA1()),
+        skip_message="Does not support HMAC-SHA1."
+    )
+    @pytest.mark.parametrize(
+        "params", [i for i in vectors if i["mode"] == b"SHA1"])
+    def test_generate_sha1(self, backend, params):
         secret = params["secret"]
         time = int(params["time"])
-        mode = params["mode"]
         totp_value = params["totp"]
 
-        algorithm = getattr(hashes, mode.decode())
-        totp = TOTP(secret, 8, algorithm(), 30, backend)
+        totp = TOTP(secret, 8, hashes.SHA1(), 30, backend)
         assert totp.generate(time) == totp_value
 
-    @pytest.mark.parametrize("params", vectors)
-    def test_verify(self, backend, params):
+    @pytest.mark.supported(
+        only_if=lambda backend: backend.hmac_supported(hashes.SHA256()),
+        skip_message="Does not support HMAC-SHA256."
+    )
+    @pytest.mark.parametrize(
+        "params", [i for i in vectors if i["mode"] == b"SHA256"])
+    def test_generate_sha256(self, backend, params):
         secret = params["secret"]
         time = int(params["time"])
-        mode = params["mode"]
         totp_value = params["totp"]
 
-        algorithm = getattr(hashes, mode.decode())
-        totp = TOTP(secret, 8, algorithm(), 30, backend)
+        totp = TOTP(secret, 8, hashes.SHA256(), 30, backend)
+        assert totp.generate(time) == totp_value
+
+    @pytest.mark.supported(
+        only_if=lambda backend: backend.hmac_supported(hashes.SHA512()),
+        skip_message="Does not support HMAC-SHA512."
+    )
+    @pytest.mark.parametrize(
+        "params", [i for i in vectors if i["mode"] == b"SHA512"])
+    def test_generate_sha512(self, backend, params):
+        secret = params["secret"]
+        time = int(params["time"])
+        totp_value = params["totp"]
+
+        totp = TOTP(secret, 8, hashes.SHA512(), 30, backend)
+        assert totp.generate(time) == totp_value
+
+    @pytest.mark.supported(
+        only_if=lambda backend: backend.hmac_supported(hashes.SHA1()),
+        skip_message="Does not support HMAC-SHA1."
+    )
+    @pytest.mark.parametrize(
+        "params", [i for i in vectors if i["mode"] == b"SHA1"])
+    def test_verify_sha1(self, backend, params):
+        secret = params["secret"]
+        time = int(params["time"])
+        totp_value = params["totp"]
+
+        totp = TOTP(secret, 8, hashes.SHA1(), 30, backend)
+
+        assert totp.verify(totp_value, time) is None
+
+    @pytest.mark.supported(
+        only_if=lambda backend: backend.hmac_supported(hashes.SHA256()),
+        skip_message="Does not support HMAC-SHA256."
+    )
+    @pytest.mark.parametrize(
+        "params", [i for i in vectors if i["mode"] == b"SHA256"])
+    def test_verify_sha256(self, backend, params):
+        secret = params["secret"]
+        time = int(params["time"])
+        totp_value = params["totp"]
+
+        totp = TOTP(secret, 8, hashes.SHA256(), 30, backend)
+
+        assert totp.verify(totp_value, time) is None
+
+    @pytest.mark.supported(
+        only_if=lambda backend: backend.hmac_supported(hashes.SHA512()),
+        skip_message="Does not support HMAC-SHA512."
+    )
+    @pytest.mark.parametrize(
+        "params", [i for i in vectors if i["mode"] == b"SHA512"])
+    def test_verify_sha512(self, backend, params):
+        secret = params["secret"]
+        time = int(params["time"])
+        totp_value = params["totp"]
+
+        totp = TOTP(secret, 8, hashes.SHA512(), 30, backend)
 
         assert totp.verify(totp_value, time) is None
 
     def test_invalid_verify(self, backend):
-            secret = b"12345678901234567890"
-            time = 59
+        secret = b"12345678901234567890"
+        time = 59
 
-            totp = TOTP(secret, 8, SHA1(), 30, backend)
+        totp = TOTP(secret, 8, hashes.SHA1(), 30, backend)
 
-            with pytest.raises(InvalidToken):
-                totp.verify(b"12345678", time)
+        with pytest.raises(InvalidToken):
+            totp.verify(b"12345678", time)
 
     def test_floating_point_time_generate(self, backend):
         secret = b"12345678901234567890"
         time = 59.1
 
-        totp = TOTP(secret, 8, SHA1(), 30, backend)
+        totp = TOTP(secret, 8, hashes.SHA1(), 30, backend)
 
         assert totp.generate(time) == b"94287082"