use subtests in aead (#5618)

1 files changed, 88 insertions, 92 deletions

diff --git a/tests/hazmat/primitives/test_aead.py b/tests/hazmat/primitives/test_aead.py
index bc7f5ab11..6787cb503 100644
--- a/tests/hazmat/primitives/test_aead.py
+++ b/tests/hazmat/primitives/test_aead.py
@@ -121,55 +121,53 @@ class TestChaCha20Poly1305(object):
         pt2 = chacha.decrypt(nonce, ct2, b"")
         assert pt1 == pt2
 
-    @pytest.mark.parametrize(
-        "vector",
-        load_vectors_from_file(
+    def test_openssl_vectors(self, subtests, backend):
+        vectors = load_vectors_from_file(
             os.path.join("ciphers", "ChaCha20Poly1305", "openssl.txt"),
             load_nist_vectors,
-        ),
-    )
-    def test_openssl_vectors(self, vector, backend):
-        key = binascii.unhexlify(vector["key"])
-        nonce = binascii.unhexlify(vector["iv"])
-        aad = binascii.unhexlify(vector["aad"])
-        tag = binascii.unhexlify(vector["tag"])
-        pt = binascii.unhexlify(vector["plaintext"])
-        ct = binascii.unhexlify(vector["ciphertext"])
-        chacha = ChaCha20Poly1305(key)
-        if vector.get("result") == b"CIPHERFINAL_ERROR":
-            with pytest.raises(InvalidTag):
-                chacha.decrypt(nonce, ct + tag, aad)
-        else:
-            computed_pt = chacha.decrypt(nonce, ct + tag, aad)
-            assert computed_pt == pt
-            computed_ct = chacha.encrypt(nonce, pt, aad)
-            assert computed_ct == ct + tag
-
-    @pytest.mark.parametrize(
-        "vector",
-        load_vectors_from_file(
+        )
+        for vector in vectors:
+            with subtests.test():
+                key = binascii.unhexlify(vector["key"])
+                nonce = binascii.unhexlify(vector["iv"])
+                aad = binascii.unhexlify(vector["aad"])
+                tag = binascii.unhexlify(vector["tag"])
+                pt = binascii.unhexlify(vector["plaintext"])
+                ct = binascii.unhexlify(vector["ciphertext"])
+                chacha = ChaCha20Poly1305(key)
+                if vector.get("result") == b"CIPHERFINAL_ERROR":
+                    with pytest.raises(InvalidTag):
+                        chacha.decrypt(nonce, ct + tag, aad)
+                else:
+                    computed_pt = chacha.decrypt(nonce, ct + tag, aad)
+                    assert computed_pt == pt
+                    computed_ct = chacha.encrypt(nonce, pt, aad)
+                    assert computed_ct == ct + tag
+
+    def test_boringssl_vectors(self, subtests, backend):
+        vectors = load_vectors_from_file(
             os.path.join("ciphers", "ChaCha20Poly1305", "boringssl.txt"),
             load_nist_vectors,
-        ),
-    )
-    def test_boringssl_vectors(self, vector, backend):
-        key = binascii.unhexlify(vector["key"])
-        nonce = binascii.unhexlify(vector["nonce"])
-        if vector["ad"].startswith(b'"'):
-            aad = vector["ad"][1:-1]
-        else:
-            aad = binascii.unhexlify(vector["ad"])
-        tag = binascii.unhexlify(vector["tag"])
-        if vector["in"].startswith(b'"'):
-            pt = vector["in"][1:-1]
-        else:
-            pt = binascii.unhexlify(vector["in"])
-        ct = binascii.unhexlify(vector["ct"].strip(b'"'))
-        chacha = ChaCha20Poly1305(key)
-        computed_pt = chacha.decrypt(nonce, ct + tag, aad)
-        assert computed_pt == pt
-        computed_ct = chacha.encrypt(nonce, pt, aad)
-        assert computed_ct == ct + tag
+        )
+        for vector in vectors:
+            with subtests.test():
+                key = binascii.unhexlify(vector["key"])
+                nonce = binascii.unhexlify(vector["nonce"])
+                if vector["ad"].startswith(b'"'):
+                    aad = vector["ad"][1:-1]
+                else:
+                    aad = binascii.unhexlify(vector["ad"])
+                tag = binascii.unhexlify(vector["tag"])
+                if vector["in"].startswith(b'"'):
+                    pt = vector["in"][1:-1]
+                else:
+                    pt = binascii.unhexlify(vector["in"])
+                ct = binascii.unhexlify(vector["ct"].strip(b'"'))
+                chacha = ChaCha20Poly1305(key)
+                computed_pt = chacha.decrypt(nonce, ct + tag, aad)
+                assert computed_pt == pt
+                computed_ct = chacha.encrypt(nonce, pt, aad)
+                assert computed_ct == ct + tag
 
     def test_buffer_protocol(self, backend):
         key = ChaCha20Poly1305.generate_key()
@@ -230,9 +228,8 @@ class TestAESCCM(object):
         with pytest.raises(ValueError):
             aesccm.encrypt(nonce[:6], pt, None)
 
-    @pytest.mark.parametrize(
-        "vector",
-        _load_all_params(
+    def test_vectors(self, subtests, backend):
+        vectors = _load_all_params(
             os.path.join("ciphers", "AES", "CCM"),
             [
                 "DVPT128.rsp",
@@ -249,22 +246,22 @@ class TestAESCCM(object):
                 "VPT256.rsp",
             ],
             load_nist_ccm_vectors,
-        ),
-    )
-    def test_vectors(self, vector, backend):
-        key = binascii.unhexlify(vector["key"])
-        nonce = binascii.unhexlify(vector["nonce"])
-        adata = binascii.unhexlify(vector["adata"])[: vector["alen"]]
-        ct = binascii.unhexlify(vector["ct"])
-        pt = binascii.unhexlify(vector["payload"])[: vector["plen"]]
-        aesccm = AESCCM(key, vector["tlen"])
-        if vector.get("fail"):
-            with pytest.raises(InvalidTag):
-                aesccm.decrypt(nonce, ct, adata)
-        else:
-            computed_pt = aesccm.decrypt(nonce, ct, adata)
-            assert computed_pt == pt
-            assert aesccm.encrypt(nonce, pt, adata) == ct
+        )
+        for vector in vectors:
+            with subtests.test():
+                key = binascii.unhexlify(vector["key"])
+                nonce = binascii.unhexlify(vector["nonce"])
+                adata = binascii.unhexlify(vector["adata"])[: vector["alen"]]
+                ct = binascii.unhexlify(vector["ct"])
+                pt = binascii.unhexlify(vector["payload"])[: vector["plen"]]
+                aesccm = AESCCM(key, vector["tlen"])
+                if vector.get("fail"):
+                    with pytest.raises(InvalidTag):
+                        aesccm.decrypt(nonce, ct, adata)
+                else:
+                    computed_pt = aesccm.decrypt(nonce, ct, adata)
+                    assert computed_pt == pt
+                    assert aesccm.encrypt(nonce, pt, adata) == ct
 
     def test_roundtrip(self, backend):
         key = AESCCM.generate_key(128)
@@ -359,7 +356,7 @@ def _load_gcm_vectors():
         ],
         load_nist_vectors,
     )
-    return [x for x in vectors if len(x["tag"]) == 32]
+    return [x for x in vectors if len(x["tag"]) == 32 and len(x["iv"]) >= 16]
 
 
 @pytest.mark.requires_backend_interface(interface=CipherBackend)
@@ -375,33 +372,32 @@ class TestAESGCM(object):
         with pytest.raises(OverflowError):
             aesgcm.encrypt(nonce, b"", FakeData())
 
-    @pytest.mark.parametrize("vector", _load_gcm_vectors())
-    def test_vectors(self, backend, vector):
-        nonce = binascii.unhexlify(vector["iv"])
-
-        if len(nonce) < 8:
-            pytest.skip("GCM does not support less than 64-bit IVs")
-
-        if backend._fips_enabled and len(nonce) != 12:
-            # Red Hat disables non-96-bit IV support as part of its FIPS
-            # patches.
-            pytest.skip("Non-96-bit IVs unsupported in FIPS mode.")
-
-        key = binascii.unhexlify(vector["key"])
-        aad = binascii.unhexlify(vector["aad"])
-        ct = binascii.unhexlify(vector["ct"])
-        pt = binascii.unhexlify(vector.get("pt", b""))
-        tag = binascii.unhexlify(vector["tag"])
-        aesgcm = AESGCM(key)
-        if vector.get("fail") is True:
-            with pytest.raises(InvalidTag):
-                aesgcm.decrypt(nonce, ct + tag, aad)
-        else:
-            computed_ct = aesgcm.encrypt(nonce, pt, aad)
-            assert computed_ct[:-16] == ct
-            assert computed_ct[-16:] == tag
-            computed_pt = aesgcm.decrypt(nonce, ct + tag, aad)
-            assert computed_pt == pt
+    def test_vectors(self, backend, subtests):
+        vectors = _load_gcm_vectors()
+        for vector in vectors:
+            with subtests.test():
+                nonce = binascii.unhexlify(vector["iv"])
+
+                if backend._fips_enabled and len(nonce) != 12:
+                    # Red Hat disables non-96-bit IV support as part of its
+                    # FIPS patches.
+                    pytest.skip("Non-96-bit IVs unsupported in FIPS mode.")
+
+                key = binascii.unhexlify(vector["key"])
+                aad = binascii.unhexlify(vector["aad"])
+                ct = binascii.unhexlify(vector["ct"])
+                pt = binascii.unhexlify(vector.get("pt", b""))
+                tag = binascii.unhexlify(vector["tag"])
+                aesgcm = AESGCM(key)
+                if vector.get("fail") is True:
+                    with pytest.raises(InvalidTag):
+                        aesgcm.decrypt(nonce, ct + tag, aad)
+                else:
+                    computed_ct = aesgcm.encrypt(nonce, pt, aad)
+                    assert computed_ct[:-16] == ct
+                    assert computed_ct[-16:] == tag
+                    computed_pt = aesgcm.decrypt(nonce, ct + tag, aad)
+                    assert computed_pt == pt
 
     @pytest.mark.parametrize(
         ("nonce", "data", "associated_data"),