feat(workos-python): Add connection to get_authorization_url (#61)

Author: rohanjadvani

tests/test_sso.py

@@ -16,7 +16,8 @@ def setup_with_client_id(self, set_api_key_and_client_id):
         self.provider = ConnectionType.GoogleOAuth
         self.customer_domain = "workos.com"
         self.redirect_uri = "https://localhost/auth/callback"
-        self.state = json.dumps({"things": "with_stuff",})
+        self.state = json.dumps({"things": "with_stuff"})
+        self.connection = "connection_123"
 
         self.sso = SSO()
 
@@ -25,7 +26,7 @@ def setup_with_project_id(self, set_api_key_and_project_id):
         self.provider = ConnectionType.GoogleOAuth
         self.customer_domain = "workos.com"
         self.redirect_uri = "https://localhost/auth/callback"
-        self.state = json.dumps({"things": "with_stuff",})
+        self.state = json.dumps({"things": "with_stuff"})
 
         self.sso = SSO()
 
@@ -99,7 +100,7 @@ def mock_connections(self):
             "listMetadata": {"before": None, "after": None},
         }
 
-    def test_authorization_url_throws_value_error_with_missing_domain_and_provider(
+    def test_authorization_url_throws_value_error_with_missing_connection_domain_and_provider(
         self, setup_with_client_id
     ):
         with pytest.raises(ValueError, match=r"Incomplete arguments.*"):
@@ -153,6 +154,25 @@ def test_authorization_url_has_expected_query_params_with_domain(
             "state": self.state,
         }
 
+    def test_authorization_url_has_expected_query_params_with_connection(
+        self, setup_with_client_id
+    ):
+        authorization_url = self.sso.get_authorization_url(
+            connection=self.connection,
+            redirect_uri=self.redirect_uri,
+            state=self.state,
+        )
+
+        parsed_url = urlparse(authorization_url)
+
+        assert dict(parse_qsl(parsed_url.query)) == {
+            "connection": self.connection,
+            "client_id": workos.client_id,
+            "redirect_uri": self.redirect_uri,
+            "response_type": RESPONSE_TYPE_CODE,
+            "state": self.state,
+        }
+
     def test_authorization_url_has_expected_query_params_with_domain_and_provider(
         self, setup_with_client_id
     ):

workos/sso.py

@@ -40,7 +40,7 @@ def request_helper(self):
         return self._request_helper
 
     def get_authorization_url(
-        self, domain=None, redirect_uri=None, state=None, provider=None
+        self, domain=None, redirect_uri=None, state=None, provider=None, connection=None
     ):
         """Generate an OAuth 2.0 authorization URL.
 
@@ -53,6 +53,7 @@ def get_authorization_url(
             state (str) - An encoded string passed to WorkOS that'd be preserved through the authentication workflow, passed
             back as a query parameter
             provider (ConnectionType) - Authentication service provider descriptor
+            connection (string) - Unique identifier for a WorkOS Connection
 
         Returns:
             str: URL to redirect a User to to begin the OAuth workflow with WorkOS
@@ -63,16 +64,18 @@ def get_authorization_url(
             "response_type": RESPONSE_TYPE_CODE,
         }
 
-        if domain is None and provider is None:
+        if domain is None and provider is None and connection is None:
             raise ValueError(
-                "Incomplete arguments. Need to specify either a 'domain' or 'provider'"
+                "Incomplete arguments. Need to specify either a 'connection', 'domain' or 'provider'"
             )
         if provider is not None:
             if not isinstance(provider, ConnectionType):
                 raise ValueError("'provider' must be of type ConnectionType")
             params["provider"] = str(provider.value)
         if domain is not None:
             params["domain"] = domain
+        if connection is not None:
+            params["connection"] = connection
 
         if state is not None:
             params["state"] = state