Merge pull request #742 from johanlundberg/bugfix_registration_policy_none

Handle registration_policy None

2 files changed, 31 insertions, 14 deletions

diff --git a/src/saml2/mdstore.py b/src/saml2/mdstore.py
index 76a981a8..96acfa1e 100644
--- a/src/saml2/mdstore.py
+++ b/src/saml2/mdstore.py
@@ -1432,22 +1432,29 @@ class MetadataStore(MetaData):
         :type entity_id: string
         :rtype: dict
         """
-        res = {
-            'registration_authority': None,
-            'registration_instant': None,
-            'registration_policy': {}
-        }
         try:
-            ext = self.__getitem__(entity_id)["extensions"]
+            ext = self.__getitem__(entity_id)
         except KeyError:
-            return res
-        for elem in ext["extension_elements"]:
-            if elem["__class__"] == classnames["mdrpi_registration_info"]:
-                res["registration_authority"] = elem["registration_authority"]
-                res["registration_instant"] = elem.get("registration_instant")
-                for policy in elem.get('registration_policy'):
-                    if policy["__class__"] == classnames["mdrpi_registration_policy"]:
-                        res['registration_policy'][policy["lang"]] = policy["text"]
+            ext = {}
+
+        ext_elems = ext.get("extensions", {}).get("extension_elements", [])
+        reg_info = next(
+            (
+                elem
+                for elem in ext_elems
+                if elem["__class__"] == classnames["mdrpi_registration_info"]
+            ),
+            {},
+        )
+        res = {
+            "registration_authority": reg_info.get("registration_authority"),
+            "registration_instant": reg_info.get("registration_instant"),
+            "registration_policy": {
+                policy["lang"]: policy["text"]
+                for policy in reg_info.get("registration_policy", [])
+                if policy["__class__"] == classnames["mdrpi_registration_policy"]
+            },
+        }
         return res
 
     def _lookup_elements_by_cls(self, root, cls):
diff --git a/tests/test_30_mdstore.py b/tests/test_30_mdstore.py
index fede1010..4dfa80f3 100644
--- a/tests/test_30_mdstore.py
+++ b/tests/test_30_mdstore.py
@@ -575,6 +575,16 @@ def test_registration_info():
            registration_info['registration_policy']['en']
 
 
+def test_registration_info_no_policy():
+    mds = MetadataStore(ATTRCONV, sec_config,
+                        disable_ssl_certificate_validation=True)
+    mds.imp(METADATACONF["13"])
+    registration_info = mds.registration_info(entity_id='https://idp.szie.hu/idp/shibboleth')
+    assert 'http://eduid.hu' == registration_info['registration_authority']
+    assert registration_info['registration_instant'] is None
+    assert registration_info['registration_policy'] == {}
+
+
 def test_extension():
     mds = MetadataStore(ATTRCONV, None)
     # use ordered dict to force expected entity to be last