Don't catch KeyError over too broad statement.

Previously KeyError in lower level of implementation,
MetadataStore.entity_categories(), was caught instead of being
propagated.

1 files changed, 5 insertions, 6 deletions

diff --git a/src/saml2/assertion.py b/src/saml2/assertion.py
index 35a11dd0..adfeecd7 100644
--- a/src/saml2/assertion.py
+++ b/src/saml2/assertion.py
@@ -273,13 +273,8 @@ def post_entity_categories(maps, **kwargs):
         required = []
 
     if kwargs["mds"]:
-        try:
+        if "sp_entity_id" in kwargs:
             ecs = kwargs["mds"].entity_categories(kwargs["sp_entity_id"])
-        except KeyError:
-            for ec_map in maps:
-                for attr in ec_map[""]:
-                    restrictions[attr] = None
-        else:
             for ec_map in maps:
                 for key, (atlist, only_required) in ec_map.items():
                     if key == "":  # always released
@@ -305,6 +300,10 @@ def post_entity_categories(maps, **kwargs):
 
                     for attr in attrs:
                         restrictions[attr] = None
+        else:
+            for ec_map in maps:
+                for attr in ec_map[""]:
+                    restrictions[attr] = None
 
     return restrictions