Try-except extract data funtcion and add unit test

aikido_zen/helpers/extract_data_from_xml_body.py

@@ -1,18 +1,26 @@
 """Exports extract_data_from_xml_body helper function"""
 
 import aikido_zen.context as ctx
+from aikido_zen.helpers.logging import logger
 
 
 def extract_data_from_xml_body(user_input, root_element):
     """Extracts all attributes from the xml and adds them to context"""
-    context = ctx.get_current_context()
-    if not isinstance(context.body, str) or user_input != context.body:
-        return
+    try:
+        context = ctx.get_current_context()
+        if (
+            not context
+            or not isinstance(context.body, str)
+            or user_input != context.body
+        ):
+            return
 
-    extracted_xml_attrs = context.xml
-    for el in root_element:
-        for k, v in el.items():
-            if not extracted_xml_attrs.get(k):
-                extracted_xml_attrs[k] = set()
-            extracted_xml_attrs[k].add(v)
-    context.set_as_current_context()
+        extracted_xml_attrs = context.xml
+        for el in root_element:
+            for k, v in el.items():
+                if not extracted_xml_attrs.get(k):
+                    extracted_xml_attrs[k] = set()
+                extracted_xml_attrs[k].add(v)
+        context.set_as_current_context()
+    except Exception as e:
+        logger.debug("Exception occured when extracting XML: %s", e)

aikido_zen/helpers/extract_data_from_xml_body_test.py

@@ -14,6 +14,17 @@ def mock_context():
     return mock_ctx
 
 
+def test_does_not_crash_when_context_none(mock_context):
+    with patch("aikido_zen.context.get_current_context", return_value=None):
+        user_input = "valid_input"
+        root_element = [
+            {"attr1": "value1", "attr2": "value2"},
+            {"attr1": "value3", "attr3": "value4"},
+        ]
+
+        extract_data_from_xml_body(user_input, root_element)
+
+
 def test_extract_data_from_xml_body_valid_input(mock_context):
     with patch("aikido_zen.context.get_current_context", return_value=mock_context):
         user_input = "valid_input"