Format code

Signed-off-by: Ivan Kanakarakis <[email protected]>

Author: c00kiemon5ter

src/satosa/backends/saml2.py

@@ -121,21 +121,18 @@ def disco_query(self, context):
         endpoints = self.sp.config.getattr("endpoints", "sp")
         return_url = endpoints["discovery_response"][0][0]
 
-        disco_url = context.get_decoration(self.KEY_SAML_DISCOVERY_SERVICE_URL)
-        disco_policy = (
-                context.get_decoration(self.KEY_SAML_DISCOVERY_SERVICE_POLICY)
-                )
-
-        if not disco_url:
-            disco_url = self.discosrv
+        disco_url = (
+            context.get_decoration(self.KEY_SAML_DISCOVERY_SERVICE_URL) or self.discosrv
+        )
+        disco_policy = context.get_decoration(self.KEY_SAML_DISCOVERY_SERVICE_POLICY)
 
-        args = {"return" : return_url}
+        args = {"return": return_url}
         if disco_policy:
             args["policy"] = disco_policy
 
         loc = self.sp.create_discovery_service_request(
-                    disco_url, self.sp.config.entityid, **args)
-
+            disco_url, self.sp.config.entityid, **args
+        )
         return SeeOther(loc)
 
     def construct_requested_authn_context(self, entity_id):

tests/satosa/backends/test_saml2.py

@@ -48,10 +48,8 @@ def assert_redirect_to_idp(self, redirect_response, idp_conf):
         assert "SAMLRequest" in parse_qs(parsed.query)
 
     def assert_redirect_to_discovery_server(
-                                            self,
-                                            redirect_response,
-                                            sp_conf,
-                                            expected_discosrv_url):
+        self, redirect_response, sp_conf, expected_discosrv_url
+    ):
         assert redirect_response.status == "303 See Other"
         parsed = urlparse(redirect_response.message)
         redirect_location = "{parsed.scheme}://{parsed.netloc}{parsed.path}".format(parsed=parsed)
@@ -108,8 +106,8 @@ def test_start_auth_defaults_to_redirecting_to_discovery_server(self, context, s
     def test_discovery_server_set_in_context(self, context, sp_conf):
         discosrv_url = 'https://my.org/saml_discovery_service'
         context.decorate(
-            SAMLBackend.KEY_SAML_DISCOVERY_SERVICE_URL,
-            discosrv_url)
+            SAMLBackend.KEY_SAML_DISCOVERY_SERVICE_URL, discosrv_url
+        )
         resp = self.samlbackend.start_auth(context, InternalData())
         self.assert_redirect_to_discovery_server(resp, sp_conf, discosrv_url)