subtest-ify keywrap (#5619)

1 files changed, 82 insertions, 80 deletions

diff --git a/tests/hazmat/primitives/test_keywrap.py b/tests/hazmat/primitives/test_keywrap.py
index 947afb3e3..c5486fe1e 100644
--- a/tests/hazmat/primitives/test_keywrap.py
+++ b/tests/hazmat/primitives/test_keywrap.py
@@ -18,14 +18,6 @@ from ...utils import load_nist_vectors
 
 @pytest.mark.requires_backend_interface(interface=CipherBackend)
 class TestAESKeyWrap(object):
-    @pytest.mark.parametrize(
-        "params",
-        _load_all_params(
-            os.path.join("keywrap", "kwtestvectors"),
-            ["KW_AE_128.txt", "KW_AE_192.txt", "KW_AE_256.txt"],
-            load_nist_vectors,
-        ),
-    )
     @pytest.mark.supported(
         only_if=lambda backend: backend.cipher_supported(
             algorithms.AES(b"\x00" * 16), modes.ECB()
@@ -33,20 +25,21 @@ class TestAESKeyWrap(object):
         skip_message="Does not support AES key wrap (RFC 3394) because AES-ECB"
         " is unsupported",
     )
-    def test_wrap(self, backend, params):
-        wrapping_key = binascii.unhexlify(params["k"])
-        key_to_wrap = binascii.unhexlify(params["p"])
-        wrapped_key = keywrap.aes_key_wrap(wrapping_key, key_to_wrap, backend)
-        assert params["c"] == binascii.hexlify(wrapped_key)
-
-    @pytest.mark.parametrize(
-        "params",
-        _load_all_params(
+    def test_wrap(self, backend, subtests):
+        params = _load_all_params(
             os.path.join("keywrap", "kwtestvectors"),
-            ["KW_AD_128.txt", "KW_AD_192.txt", "KW_AD_256.txt"],
+            ["KW_AE_128.txt", "KW_AE_192.txt", "KW_AE_256.txt"],
             load_nist_vectors,
-        ),
-    )
+        )
+        for param in params:
+            with subtests.test():
+                wrapping_key = binascii.unhexlify(param["k"])
+                key_to_wrap = binascii.unhexlify(param["p"])
+                wrapped_key = keywrap.aes_key_wrap(
+                    wrapping_key, key_to_wrap, backend
+                )
+                assert param["c"] == binascii.hexlify(wrapped_key)
+
     @pytest.mark.supported(
         only_if=lambda backend: backend.cipher_supported(
             algorithms.AES(b"\x00" * 16), modes.ECB()
@@ -54,17 +47,26 @@ class TestAESKeyWrap(object):
         skip_message="Does not support AES key wrap (RFC 3394) because AES-ECB"
         " is unsupported",
     )
-    def test_unwrap(self, backend, params):
-        wrapping_key = binascii.unhexlify(params["k"])
-        wrapped_key = binascii.unhexlify(params["c"])
-        if params.get("fail") is True:
-            with pytest.raises(keywrap.InvalidUnwrap):
-                keywrap.aes_key_unwrap(wrapping_key, wrapped_key, backend)
-        else:
-            unwrapped_key = keywrap.aes_key_unwrap(
-                wrapping_key, wrapped_key, backend
-            )
-            assert params["p"] == binascii.hexlify(unwrapped_key)
+    def test_unwrap(self, backend, subtests):
+        params = _load_all_params(
+            os.path.join("keywrap", "kwtestvectors"),
+            ["KW_AD_128.txt", "KW_AD_192.txt", "KW_AD_256.txt"],
+            load_nist_vectors,
+        )
+        for param in params:
+            with subtests.test():
+                wrapping_key = binascii.unhexlify(param["k"])
+                wrapped_key = binascii.unhexlify(param["c"])
+                if param.get("fail") is True:
+                    with pytest.raises(keywrap.InvalidUnwrap):
+                        keywrap.aes_key_unwrap(
+                            wrapping_key, wrapped_key, backend
+                        )
+                else:
+                    unwrapped_key = keywrap.aes_key_unwrap(
+                        wrapping_key, wrapped_key, backend
+                    )
+                    assert param["p"] == binascii.hexlify(unwrapped_key)
 
     @pytest.mark.supported(
         only_if=lambda backend: backend.cipher_supported(
@@ -124,67 +126,67 @@ class TestAESKeyWrap(object):
 )
 @pytest.mark.requires_backend_interface(interface=CipherBackend)
 class TestAESKeyWrapWithPadding(object):
-    @pytest.mark.parametrize(
-        "params",
-        _load_all_params(
+    def test_wrap(self, backend, subtests):
+        params = _load_all_params(
             os.path.join("keywrap", "kwtestvectors"),
             ["KWP_AE_128.txt", "KWP_AE_192.txt", "KWP_AE_256.txt"],
             load_nist_vectors,
-        ),
-    )
-    def test_wrap(self, backend, params):
-        wrapping_key = binascii.unhexlify(params["k"])
-        key_to_wrap = binascii.unhexlify(params["p"])
-        wrapped_key = keywrap.aes_key_wrap_with_padding(
-            wrapping_key, key_to_wrap, backend
         )
-        assert params["c"] == binascii.hexlify(wrapped_key)
+        for param in params:
+            with subtests.test():
+                wrapping_key = binascii.unhexlify(param["k"])
+                key_to_wrap = binascii.unhexlify(param["p"])
+                wrapped_key = keywrap.aes_key_wrap_with_padding(
+                    wrapping_key, key_to_wrap, backend
+                )
+                assert param["c"] == binascii.hexlify(wrapped_key)
 
-    @pytest.mark.parametrize(
-        "params",
-        _load_all_params("keywrap", ["kwp_botan.txt"], load_nist_vectors),
-    )
-    def test_wrap_additional_vectors(self, backend, params):
-        wrapping_key = binascii.unhexlify(params["key"])
-        key_to_wrap = binascii.unhexlify(params["input"])
-        wrapped_key = keywrap.aes_key_wrap_with_padding(
-            wrapping_key, key_to_wrap, backend
+    def test_wrap_additional_vectors(self, backend, subtests):
+        params = _load_all_params(
+            "keywrap", ["kwp_botan.txt"], load_nist_vectors
         )
-        assert wrapped_key == binascii.unhexlify(params["output"])
+        for param in params:
+            with subtests.test():
+                wrapping_key = binascii.unhexlify(param["key"])
+                key_to_wrap = binascii.unhexlify(param["input"])
+                wrapped_key = keywrap.aes_key_wrap_with_padding(
+                    wrapping_key, key_to_wrap, backend
+                )
+                assert wrapped_key == binascii.unhexlify(param["output"])
 
-    @pytest.mark.parametrize(
-        "params",
-        _load_all_params(
+    def test_unwrap(self, backend, subtests):
+        params = _load_all_params(
             os.path.join("keywrap", "kwtestvectors"),
             ["KWP_AD_128.txt", "KWP_AD_192.txt", "KWP_AD_256.txt"],
             load_nist_vectors,
-        ),
-    )
-    def test_unwrap(self, backend, params):
-        wrapping_key = binascii.unhexlify(params["k"])
-        wrapped_key = binascii.unhexlify(params["c"])
-        if params.get("fail") is True:
-            with pytest.raises(keywrap.InvalidUnwrap):
-                keywrap.aes_key_unwrap_with_padding(
+        )
+        for param in params:
+            with subtests.test():
+                wrapping_key = binascii.unhexlify(param["k"])
+                wrapped_key = binascii.unhexlify(param["c"])
+                if param.get("fail") is True:
+                    with pytest.raises(keywrap.InvalidUnwrap):
+                        keywrap.aes_key_unwrap_with_padding(
+                            wrapping_key, wrapped_key, backend
+                        )
+                else:
+                    unwrapped_key = keywrap.aes_key_unwrap_with_padding(
+                        wrapping_key, wrapped_key, backend
+                    )
+                    assert param["p"] == binascii.hexlify(unwrapped_key)
+
+    def test_unwrap_additional_vectors(self, backend, subtests):
+        params = _load_all_params(
+            "keywrap", ["kwp_botan.txt"], load_nist_vectors
+        )
+        for param in params:
+            with subtests.test():
+                wrapping_key = binascii.unhexlify(param["key"])
+                wrapped_key = binascii.unhexlify(param["output"])
+                unwrapped_key = keywrap.aes_key_unwrap_with_padding(
                     wrapping_key, wrapped_key, backend
                 )
-        else:
-            unwrapped_key = keywrap.aes_key_unwrap_with_padding(
-                wrapping_key, wrapped_key, backend
-            )
-            assert params["p"] == binascii.hexlify(unwrapped_key)
-
-    @pytest.mark.parametrize(
-        "params",
-        _load_all_params("keywrap", ["kwp_botan.txt"], load_nist_vectors),
-    )
-    def test_unwrap_additional_vectors(self, backend, params):
-        wrapping_key = binascii.unhexlify(params["key"])
-        wrapped_key = binascii.unhexlify(params["output"])
-        unwrapped_key = keywrap.aes_key_unwrap_with_padding(
-            wrapping_key, wrapped_key, backend
-        )
-        assert unwrapped_key == binascii.unhexlify(params["input"])
+                assert unwrapped_key == binascii.unhexlify(param["input"])
 
     def test_unwrap_invalid_wrapped_key_length(self, backend):
         # Keys to unwrap must be at least 16 bytes