Fix addSign test

Author: pitbulk

tests/src/OneLogin/saml2_tests/utils_test.py

@@ -813,51 +813,29 @@ def testAddSign(self):
 
         xml_authn = b64decode(self.file_contents(join(self.data_path, 'requests', 'authn_request.xml.base64')))
         xml_authn_signed = OneLogin_Saml2_Utils.add_sign(xml_authn, key, cert)
-        self.assertIn('<ds:SignatureValue>', xml_authn_signed)
-
-        res = parseString(xml_authn_signed)
-        ds_signature = res.firstChild.firstChild.nextSibling.nextSibling.nextSibling
-        self.assertEqual('ds:Signature', ds_signature.tagName)
 
         xml_authn_dom = parseString(xml_authn)
         xml_authn_signed_2 = OneLogin_Saml2_Utils.add_sign(xml_authn_dom, key, cert)
-        self.assertIn('<ds:SignatureValue>', xml_authn_signed_2)
-        res_2 = parseString(xml_authn_signed_2)
-        ds_signature_2 = res_2.firstChild.firstChild.nextSibling.nextSibling.nextSibling
-        self.assertEqual('ds:Signature', ds_signature_2.tagName)
-
         xml_authn_signed_3 = OneLogin_Saml2_Utils.add_sign(xml_authn_dom.firstChild, key, cert)
-        self.assertIn('<ds:SignatureValue>', xml_authn_signed_3)
-        res_3 = parseString(xml_authn_signed_3)
-        ds_signature_3 = res_3.firstChild.firstChild.nextSibling.nextSibling.nextSibling
-        self.assertEqual('ds:Signature', ds_signature_3.tagName)
 
         xml_authn_etree = etree.fromstring(xml_authn)
         xml_authn_signed_4 = OneLogin_Saml2_Utils.add_sign(xml_authn_etree, key, cert)
-        self.assertIn('<ds:SignatureValue>', xml_authn_signed_4)
-        res_4 = parseString(xml_authn_signed_4)
-        ds_signature_4 = res_4.firstChild.firstChild.nextSibling.nextSibling.nextSibling
-        self.assertEqual('ds:Signature', ds_signature_4.tagName)
-
-        xml_authn_signed_5 = OneLogin_Saml2_Utils.add_sign(xml_authn_etree, key, cert)
-        self.assertIn('<ds:SignatureValue>', xml_authn_signed_5)
-        res_5 = parseString(xml_authn_signed_5)
-        ds_signature_5 = res_5.firstChild.firstChild.nextSibling.nextSibling.nextSibling
-        self.assertEqual('ds:Signature', ds_signature_5.tagName)
 
         xml_logout_req = b64decode(self.file_contents(join(self.data_path, 'logout_requests', 'logout_request.xml.base64')))
         xml_logout_req_signed = OneLogin_Saml2_Utils.add_sign(xml_logout_req, key, cert)
-        self.assertIn('<ds:SignatureValue>', xml_logout_req_signed)
-        res_6 = parseString(xml_logout_req_signed)
-        ds_signature_6 = res_6.firstChild.firstChild.nextSibling.nextSibling.nextSibling
-        self.assertEqual('ds:Signature', ds_signature_6.tagName)
 
         xml_logout_res = b64decode(self.file_contents(join(self.data_path, 'logout_responses', 'logout_response.xml.base64')))
         xml_logout_res_signed = OneLogin_Saml2_Utils.add_sign(xml_logout_res, key, cert)
-        self.assertIn('<ds:SignatureValue>', xml_logout_res_signed)
-        res_7 = parseString(xml_logout_res_signed)
-        ds_signature_7 = res_7.firstChild.firstChild.nextSibling.nextSibling.nextSibling
-        self.assertEqual('ds:Signature', ds_signature_7.tagName)
+
+        xml_signed_saml_messages = [xml_authn_signed, xml_authn_signed_2, xml_authn_signed_3, xml_authn_signed_4, xml_logout_req_signed, xml_logout_res_signed]
+
+        for xml_signed_saml_message in xml_signed_saml_messages:
+            self.assertIn('<ds:SignatureValue>', xml_signed_saml_message)
+            signed_saml_message = parseString(xml_signed_saml_message)
+            ds_signature = signed_saml_message.firstChild.firstChild.nextSibling.nextSibling
+            if ds_signature.tagName != 'ds:Signature':
+                ds_signature = ds_signature.nextSibling
+            self.assertEqual('ds:Signature', ds_signature.tagName)
 
         xml_metadata = self.file_contents(join(self.data_path, 'metadata', 'metadata_settings1.xml'))
         xml_metadata_signed = OneLogin_Saml2_Utils.add_sign(xml_metadata, key, cert)