From 64093144c914d4ab832cd3e53d137a39155260ca Mon Sep 17 00:00:00 2001
From: Paul Asjes <108872335+PaulAsjes@users.noreply.github.com>
Date: Fri, 17 May 2024 15:45:54 +0200
Subject: [PATCH] Add JWKS and logout methods (#265)

* Add jwks method

* Add logout method

* ran black

* I'm getting mixed messages here black

* Ran black again, this time with the same version as CI
---
 tests/test_user_management.py | 15 +++++++++++++++
 workos/user_management.py     | 25 +++++++++++++++++++++++++
 2 files changed, 40 insertions(+)

diff --git a/tests/test_user_management.py b/tests/test_user_management.py
index 2a748db9..82588272 100644
--- a/tests/test_user_management.py
+++ b/tests/test_user_management.py
@@ -794,6 +794,21 @@ def test_authenticate_with_refresh_token(
         assert request["json"]["client_secret"] == "sk_test"
         assert request["json"]["grant_type"] == "refresh_token"
 
+    def test_get_jwks_url(self):
+        expected = "%s/sso/jwks/%s" % (workos.base_api_url, workos.client_id)
+        result = self.user_management.get_jwks_url()
+
+        assert expected == result
+
+    def test_get_logout_url(self):
+        expected = "%s/user_management/sessions/logout?session_id=%s" % (
+            workos.base_api_url,
+            "session_123",
+        )
+        result = self.user_management.get_logout_url("session_123")
+
+        assert expected == result
+
     def test_send_password_reset_email(self, capture_and_mock_request):
         email = "marcelina@foo-corp.com"
         password_reset_url = "https://foo-corp.com/reset-password"
diff --git a/workos/user_management.py b/workos/user_management.py
index d128fcc5..c20b9dce 100644
--- a/workos/user_management.py
+++ b/workos/user_management.py
@@ -761,6 +761,30 @@ def authenticate_with_refresh_token(
             response
         ).to_dict()
 
+    def get_jwks_url(self):
+        """Get the public key that is used for verifying access tokens.
+
+        Returns:
+            (str): The public JWKS URL.
+        """
+
+        return "%s/sso/jwks/%s" % (workos.base_api_url, workos.client_id)
+
+    def get_logout_url(self, session_id):
+        """Get the URL for ending the session and redirecting the user
+
+        Kwargs:
+            session_id (str): The ID of the user's session
+
+        Returns:
+            (str): URL to redirect the user to to end the session.
+        """
+
+        return "%s/user_management/sessions/logout?session_id=%s" % (
+            workos.base_api_url,
+            session_id,
+        )
+
     def send_password_reset_email(
         self,
         email,
@@ -1001,6 +1025,7 @@ def enroll_auth_factor(
         ] = WorkOSAuthenticationFactorTotp.construct_from_response(
             response["authentication_factor"]
         ).to_dict()
+
         factor_and_challenge[
             "authentication_challenge"
         ] = WorkOSChallenge.construct_from_response(