Return is a protected python word so it can not be used as a parameter name in a method definition.

These fixes are there to allow for the usage or return_url or return as parameter name.

5 files changed, 45 insertions, 20 deletions

diff --git a/src/saml2/client_base.py b/src/saml2/client_base.py
index a17f9783..e1550c96 100644
--- a/src/saml2/client_base.py
+++ b/src/saml2/client_base.py
@@ -707,12 +707,20 @@ class Base(Entity):
         :return: A URL
         """
         args = {"entityID": entity_id}
-        for key in ["return", "policy", "returnIDParam"]:
+        for key in ["policy", "returnIDParam"]:
             try:
                 args[key] = kwargs[key]
             except KeyError:
                 pass
 
+        try:
+            args["return"] = kwargs["return_url"]
+        except KeyError:
+            try:
+                args["return"] = kwargs["return"]
+            except KeyError:
+                pass
+
         if "isPassive" in kwargs:
             if kwargs["isPassive"]:
                 args["isPassive"] = "true"
diff --git a/src/saml2/discovery.py b/src/saml2/discovery.py
index 26b9de90..160b3940 100644
--- a/src/saml2/discovery.py
+++ b/src/saml2/discovery.py
@@ -62,9 +62,12 @@ class DiscoveryServer(Entity):
 
     # -------------------------------------------------------------------------
 
-    def create_discovery_service_response(self, return_url,
+    def create_discovery_service_response(self, return_url=None,
                                           returnIDParam="entityID",
-                                          entity_id=None):
+                                          entity_id=None, **kwargs):
+        if return_url is None:
+            return_url = kwargs["return"]
+            
         if entity_id:
             qp = urlencode({returnIDParam: entity_id})
 
diff --git a/tests/disco_conf.py b/tests/disco_conf.py
index 1f7554d7..253dc22e 100644
--- a/tests/disco_conf.py
+++ b/tests/disco_conf.py
@@ -1,23 +1,24 @@
 from saml2.extension.idpdisc import BINDING_DISCO
 
 from pathutils import full_path
+from pathutils import xmlsec_path
 
 BASE = "http://localhost:8088"
 
 CONFIG = {
-    "entityid" : "%s/disco.xml" % BASE,
-    "name" : "Rolands Discoserver",
+    "entityid": "%s/disco.xml" % BASE,
+    "name": "Rolands Discoserver",
     "service": {
         "ds": {
-            "endpoints" : {
+            "endpoints": {
                 "disco_service": [
                     ("%s/disco" % BASE, BINDING_DISCO),
                 ]
             },
         },
     },
-    "debug" : 1,
-    "xmlsec_binary" : None,
+    "debug": 1,
+    "xmlsec_binary": xmlsec_path,
     "metadata": {
         "local": [full_path("servera.xml")],
     },
diff --git a/tests/servera_conf.py b/tests/servera_conf.py
index 17c5f946..8d1c58c0 100644
--- a/tests/servera_conf.py
+++ b/tests/servera_conf.py
@@ -8,6 +8,7 @@ from saml2.saml import NAMEID_FORMAT_TRANSIENT
 from saml2.saml import NAMEID_FORMAT_PERSISTENT
 
 from pathutils import full_path
+from pathutils import xmlsec_path
 
 BASE = "http://lingon.catalogix.se:8087"
 
@@ -49,7 +50,7 @@ CONFIG = {
     "key_file": full_path("test.key"),
     "cert_file": full_path("test.pem"),
     "ca_certs": full_path("cacerts.txt"),
-    "xmlsec_binary": None,
+    "xmlsec_binary": xmlsec_path,
     "metadata": {
         "local": [full_path("idp_all.xml"), full_path("vo_metadata.xml")],
     },
diff --git a/tests/test_69_discovery.py b/tests/test_69_discovery.py
index 9c321259..7a14d3ea 100644
--- a/tests/test_69_discovery.py
+++ b/tests/test_69_discovery.py
@@ -5,14 +5,17 @@ from pathutils import dotname
 
 __author__ = 'rolandh'
 
-def _eq(l1,l2):
+
+def _eq(l1, l2):
     return set(l1) == set(l2)
 
+
 def test_verify():
     ds = DiscoveryServer(config_file=dotname("disco_conf"))
     assert ds
     assert ds.verify_sp_in_metadata("urn:mace:example.com:saml:roland:sp")
 
+
 def test_construct_0():
     sp = Saml2Client(config_file=dotname("servera_conf"))
     url = sp.create_discovery_service_request("http://example.com/saml/disco",
@@ -20,6 +23,8 @@ def test_construct_0():
 
     assert url == "http://example.com/saml/disco?entityID=https%3A%2F%2Fexample.com%2Fsaml%2Fsp.xml"
 
+
+
 def test_construct_1():
     sp = Saml2Client(config_file=dotname("servera_conf"))
     url = sp.create_discovery_service_request("http://example.com/saml/disco",
@@ -27,21 +32,24 @@ def test_construct_1():
 
     assert url == "http://example.com/saml/disco?entityID=https%3A%2F%2Fexample.com%2Fsaml%2Fsp.xml"
 
+
 def test_construct_deconstruct_request():
     sp = Saml2Client(config_file=dotname("servera_conf"))
-    url = sp.create_discovery_service_request("http://example.com/saml/disco",
-                                              "https://example.com/saml/sp.xml",
-                                              is_passive=True,
-                                              returnIDParam="foo",
-                                              return_url="https://example.com/saml/sp/disc")
+    url = sp.create_discovery_service_request(
+        "http://example.com/saml/disco",
+        "https://example.com/saml/sp.xml",
+        is_passive=True,
+        returnIDParam="foo",
+        return_url="https://example.com/saml/sp/disc")
 
     print url
 
     ds = DiscoveryServer(config_file=dotname("disco_conf"))
     dsr = ds.parse_discovery_service_request(url)
     # policy is added by the parsing and verifying method
-    assert _eq(dsr.keys(),["return_url", "entityID", "returnIDParam",
-                           "isPassive", "policy"])
+    assert _eq(dsr.keys(), ["return", "entityID", "returnIDParam",
+                            "isPassive", "policy"])
+
 
 def test_construct_deconstruct_response():
     sp = Saml2Client(config_file=dotname("servera_conf"))
@@ -52,10 +60,14 @@ def test_construct_deconstruct_response():
                                               return_url="https://example.com/saml/sp/disc")
     ds = DiscoveryServer(config_file=dotname("disco_conf"))
     dsr = ds.parse_discovery_service_request(url)
-    args = dict([(key, dsr[key]) for key in ["returnIDParam", "return_url"]])
+    args = dict([(key, dsr[key]) for key in ["returnIDParam", "return"]])
     url = ds.create_discovery_service_response(
-                                  entity_id="https://example.com/saml/idp.xml",
-                                  **args)
+        entity_id="https://example.com/saml/idp.xml",
+        **args)
 
     idp_id = sp.parse_discovery_service_response(url, returnIDParam="foo")
     assert idp_id == "https://example.com/saml/idp.xml"
+
+
+if __name__ == "__main__":
+    test_construct_deconstruct_response()
\ No newline at end of file