Allow optional provider parameter when building authorization URL (#10)

* Allow optional provider parameter when building authorization URL

* Update README

* Exception -> ValueError

* Add connection type enum

* Fix formatting

* Send provider and domain

* Update documentation

Author: Rohan Jadvani

README.md

@@ -30,8 +30,8 @@ from workos import client
 
 # URL to redirect a User to to initiate the WorkOS OAuth 2.0 workflow
 client.sso.get_authorization_url(
-    'customer-domain.com',
-    'my-domain.com/auth/callback',
+    domain='customer-domain.com',
+    redirect_uri='my-domain.com/auth/callback',
     state={
         'stuff': 'from_the_original_request',
         'more_things': 'ill_get_it_all_back_when_oauth_is_complete',

tests/test_sso.py

@@ -5,12 +5,14 @@
 
 import workos
 from workos.sso import SSO
+from workos.utils.connection_types import ConnectionType
 from workos.utils.request import RESPONSE_TYPE_CODE
 
 
 class TestSSO(object):
     @pytest.fixture(autouse=True)
     def setup(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 = {
@@ -30,15 +32,67 @@ def mock_profile(self):
             "idp_id": "00u1klkowm8EGah2H357",
         }
 
-    def test_authorization_url_has_expected_query_params(self):
+    def test_authorization_url_throws_value_error_with_missing_domain_and_provider(
+        self,
+    ):
+        with pytest.raises(ValueError, match=r"Incomplete arguments.*"):
+            self.sso.get_authorization_url(
+                redirect_uri=self.redirect_uri, state=self.state
+            )
+
+    def test_authorization_url_throws_value_error_with_incorrect_provider_type(self):
+        with pytest.raises(
+            ValueError, match="'provider' must be of type ConnectionType"
+        ):
+            self.sso.get_authorization_url(
+                provider="foo", redirect_uri=self.redirect_uri, state=self.state
+            )
+
+    def test_authorization_url_has_expected_query_params_with_provider(self):
+        authorization_url = self.sso.get_authorization_url(
+            provider=self.provider, redirect_uri=self.redirect_uri, state=self.state
+        )
+
+        parsed_url = urlparse(authorization_url)
+
+        assert dict(parse_qsl(parsed_url.query)) == {
+            "provider": str(self.provider),
+            "client_id": workos.project_id,
+            "redirect_uri": self.redirect_uri,
+            "response_type": RESPONSE_TYPE_CODE,
+            "state": json.dumps(self.state),
+        }
+
+    def test_authorization_url_has_expected_query_params_with_domain(self):
+        authorization_url = self.sso.get_authorization_url(
+            domain=self.customer_domain,
+            redirect_uri=self.redirect_uri,
+            state=self.state,
+        )
+
+        parsed_url = urlparse(authorization_url)
+
+        assert dict(parse_qsl(parsed_url.query)) == {
+            "domain": self.customer_domain,
+            "client_id": workos.project_id,
+            "redirect_uri": self.redirect_uri,
+            "response_type": RESPONSE_TYPE_CODE,
+            "state": json.dumps(self.state),
+        }
+
+    def test_authorization_url_has_expected_query_params_with_domain_and_provider(self):
         authorization_url = self.sso.get_authorization_url(
-            self.customer_domain, self.redirect_uri, state=self.state
+            domain=self.customer_domain,
+            provider=self.provider,
+            redirect_uri=self.redirect_uri,
+            state=self.state,
         )
 
         parsed_url = urlparse(authorization_url)
 
         assert dict(parse_qsl(parsed_url.query)) == {
             "domain": self.customer_domain,
+            "provider": str(self.provider),
             "client_id": workos.project_id,
             "redirect_uri": self.redirect_uri,
             "response_type": RESPONSE_TYPE_CODE,

workos/sso.py

@@ -5,6 +5,7 @@
 import workos
 from workos.exceptions import ConfigurationException
 from workos.resources.sso import WorkOSProfile
+from workos.utils.connection_types import ConnectionType
 from workos.utils.request import RequestHelper, RESPONSE_TYPE_CODE, REQUEST_METHOD_POST
 from workos.utils.validation import validate_api_key_and_project_id
 
@@ -27,27 +28,41 @@ def request_helper(self):
             self._request_helper = RequestHelper()
         return self._request_helper
 
-    def get_authorization_url(self, domain, redirect_uri, state=None):
+    def get_authorization_url(
+        self, domain=None, redirect_uri=None, state=None, provider=None
+    ):
         """Generate an OAuth 2.0 authorization URL.
 
         The URL generated will redirect a User to the Identity Provider configured through
         WorkOS.
 
-        Args:
+        Kwargs:
             domain (str) - The domain a user is associated with, as configured on WorkOS
             redirect_uri (str) - A valid redirect URI, as specified on WorkOS
             state (dict) - A dict passed to WorkOS, that'd be preserved through the authentication workflow, passed
             back as a query parameter
+            provider (str) - Authentication service provider descriptor
 
         Returns:
             str: URL to redirect a User to to begin the OAuth workflow with WorkOS
         """
         params = {
-            "domain": domain,
             "client_id": workos.project_id,
             "redirect_uri": redirect_uri,
             "response_type": RESPONSE_TYPE_CODE,
         }
+
+        if domain is None and provider is None:
+            raise ValueError(
+                "Incomplete arguments. Need to specify either a '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)
+        if domain is not None:
+            params["domain"] = domain
+
         if state is not None:
             params["state"] = json.dumps(state)
 

workos/utils/connection_types.py

@@ -0,0 +1,8 @@
+from enum import Enum
+
+
+class ConnectionType(Enum):
+    ADFSSAML = "ADFSSAML"
+    AzureSAML = "AzureSAML"
+    GoogleOAuth = "GoogleOAuth"
+    OktaSAML = "OktaSAML"