Merge pull request #2267 from queenp/iss2255

added get_extension_for_class #2255

1 files changed, 25 insertions, 0 deletions

diff --git a/tests/test_x509_ext.py b/tests/test_x509_ext.py
index 2c5438a9..85373973 100644
--- a/tests/test_x509_ext.py
+++ b/tests/test_x509_ext.py
@@ -832,6 +832,31 @@ class TestExtensions(object):
         extensions = cert.extensions
         assert len(extensions) == 0
 
+    def test_no_extensions_get_for_class(self, backend):
+        cert = _load_cert(
+            os.path.join(
+                "x509", "cryptography.io.pem"
+            ),
+            x509.load_pem_x509_certificate,
+            backend
+        )
+        exts = cert.extensions
+        with pytest.raises(x509.ExtensionNotFound) as exc:
+            exts.get_extension_for_class(x509.IssuerAlternativeName)
+        assert exc.value.oid == ExtensionOID.ISSUER_ALTERNATIVE_NAME
+
+    def test_one_extension_get_for_class(self, backend):
+        cert = _load_cert(
+            os.path.join(
+                "x509", "custom", "basic_constraints_not_critical.pem"
+            ),
+            x509.load_pem_x509_certificate,
+            backend
+        )
+        ext = cert.extensions.get_extension_for_class(x509.BasicConstraints)
+        assert ext is not None
+        assert isinstance(ext.value, x509.BasicConstraints)
+
 
 @pytest.mark.requires_backend_interface(interface=RSABackend)
 @pytest.mark.requires_backend_interface(interface=X509Backend)