Skip to content

Commit 13bb073

Browse files
PaulAsjesPaulAsjes
committed
Mock JWKS call to speed up tests
1 parent 87fadcb commit 13bb073

File tree

2 files changed

+66
-74
lines changed

2 files changed

+66
-74
lines changed

‎tests/conftest.py‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,9 @@
2424
from workos.utils.http_client import AsyncHTTPClient, HTTPClient, SyncHTTPClient
2525
from workos.utils.request_helper import DEFAULT_LIST_RESPONSE_LIMIT
2626

27+
from jwt import PyJWKClient
28+
from unittest.mock import Mock, patch
29+
from functools import wraps
2730

2831
def _get_test_client_setup(
2932
http_client_class_name: str,
@@ -302,3 +305,17 @@ def inner(
302305
assert request_kwargs["params"][param] == params[param]
303306

304307
return inner
308+
309+
def with_jwks_mock(func):
310+
@wraps(func)
311+
def wrapper(*args, **kwargs):
312+
# Create mock JWKS client
313+
mock_jwks = Mock(spec=PyJWKClient)
314+
mock_signing_key = Mock()
315+
mock_signing_key.key = kwargs['TEST_CONSTANTS']["PUBLIC_KEY"]
316+
mock_jwks.get_signing_key_from_jwt.return_value = mock_signing_key
317+
318+
# Apply the mock
319+
with patch('workos.session.PyJWKClient', return_value=mock_jwks):
320+
return func(*args, **kwargs)
321+
return wrapper

‎tests/test_session.py‎

Lines changed: 49 additions & 74 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
from jwt import PyJWKClient
55
from datetime import datetime, timezone
66

7+
from tests.conftest import with_jwks_mock
78
from workos.session import SessionModule
89
from workos.types.user_management.authentication_response import RefreshTokenAuthenticationResponse
910
from workos.types.user_management.session import (
@@ -60,34 +61,13 @@ def TEST_CONSTANTS():
6061
}
6162

6263
@pytest.fixture
63-
def mock_user_management(TEST_CONSTANTS):
64-
mock_jwks = Mock(spec=PyJWKClient)
65-
mock_jwk_set = Mock()
66-
mock_jwk_set.keys = [Mock(Algorithm="RS256")]
67-
mock_jwks.get_jwk_set.return_value = mock_jwk_set
68-
69-
70-
64+
def mock_user_management():
7165
mock = Mock()
7266
mock.get_jwks_url.return_value = "https://api.workos.com/user_management/sso/jwks/client_123"
73-
mock.authenticate_with_refresh_token.return_value = RefreshTokenAuthenticationResponse(
74-
**{
75-
"access_token": "access_token_123",
76-
"refresh_token": "refresh_token_123",
77-
"user": {
78-
"object": "user",
79-
"id": TEST_CONSTANTS["USER_ID"],
80-
"email": "[email protected]",
81-
"first_name": "Test",
82-
"last_name": "User",
83-
"email_verified": True,
84-
"created_at": TEST_CONSTANTS["CURRENT_TIMESTAMP"],
85-
"updated_at": TEST_CONSTANTS["CURRENT_TIMESTAMP"],
86-
}
87-
}
88-
)
67+
8968
return mock
9069

70+
@with_jwks_mock
9171
def test_initialize_session_module(TEST_CONSTANTS, mock_user_management):
9272
session = SessionModule(
9373
user_management=mock_user_management,
@@ -99,15 +79,17 @@ def test_initialize_session_module(TEST_CONSTANTS, mock_user_management):
9979
assert session.client_id == TEST_CONSTANTS["CLIENT_ID"]
10080
assert session.cookie_password is not None
10181

102-
def test_initialize_without_cookie_password(mock_user_management):
82+
@with_jwks_mock
83+
def test_initialize_without_cookie_password(TEST_CONSTANTS, mock_user_management):
10384
with pytest.raises(ValueError, match="cookie_password is required"):
10485
SessionModule(
10586
user_management=mock_user_management,
106-
client_id="client_123",
107-
session_data="session_data",
87+
client_id=TEST_CONSTANTS["CLIENT_ID"],
88+
session_data=TEST_CONSTANTS["SESSION_DATA"],
10889
cookie_password=""
10990
)
11091

92+
@with_jwks_mock
11193
def test_authenticate_no_session_cookie_provided(TEST_CONSTANTS, mock_user_management):
11294
session = SessionModule(
11395
user_management=mock_user_management,
@@ -120,6 +102,7 @@ def test_authenticate_no_session_cookie_provided(TEST_CONSTANTS, mock_user_manag
120102

121103
assert response.reason == AuthenticateWithSessionCookieFailureReason.NO_SESSION_COOKIE_PROVIDED
122104

105+
@with_jwks_mock
123106
def test_authenticate_invalid_session_cookie(TEST_CONSTANTS, mock_user_management):
124107
session = SessionModule(
125108
user_management=mock_user_management,
@@ -132,6 +115,7 @@ def test_authenticate_invalid_session_cookie(TEST_CONSTANTS, mock_user_managemen
132115

133116
assert response.reason == AuthenticateWithSessionCookieFailureReason.INVALID_SESSION_COOKIE
134117

118+
@with_jwks_mock
135119
def test_authenticate_invalid_jwt(TEST_CONSTANTS, mock_user_management):
136120
invalid_session_data = SessionModule.seal_data({ "access_token": "invalid_session_data" }, TEST_CONSTANTS["COOKIE_PASSWORD"])
137121
session = SessionModule(
@@ -145,6 +129,7 @@ def test_authenticate_invalid_jwt(TEST_CONSTANTS, mock_user_management):
145129

146130
assert response.reason == AuthenticateWithSessionCookieFailureReason.INVALID_JWT
147131

132+
@with_jwks_mock
148133
def test_authenticate_success(TEST_CONSTANTS, mock_user_management):
149134
session = SessionModule(
150135
user_management=mock_user_management,
@@ -225,6 +210,7 @@ def test_authenticate_success(TEST_CONSTANTS, mock_user_management):
225210
assert response.user.id == TEST_CONSTANTS["USER_ID"]
226211
assert response.impersonator is None
227212

213+
@with_jwks_mock
228214
def test_refresh_invalid_session_cookie(TEST_CONSTANTS, mock_user_management):
229215
session = SessionModule(
230216
user_management=mock_user_management,
@@ -238,13 +224,8 @@ def test_refresh_invalid_session_cookie(TEST_CONSTANTS, mock_user_management):
238224
assert isinstance(response, RefreshWithSessionCookieErrorResponse)
239225
assert response.reason == AuthenticateWithSessionCookieFailureReason.INVALID_SESSION_COOKIE
240226

227+
@with_jwks_mock
241228
def test_refresh_success(TEST_CONSTANTS, mock_user_management):
242-
# Create mock JWKS client
243-
mock_jwks = Mock(spec=PyJWKClient)
244-
mock_signing_key = Mock()
245-
mock_signing_key.key = TEST_CONSTANTS["PUBLIC_KEY"]
246-
mock_jwks.get_signing_key_from_jwt.return_value = mock_signing_key
247-
248229
test_user = {
249230
"object": "user",
250231
"id": TEST_CONSTANTS["USER_ID"],
@@ -272,51 +253,45 @@ def test_refresh_success(TEST_CONSTANTS, mock_user_management):
272253
**mock_response
273254
)
274255

256+
session = SessionModule(
257+
user_management=mock_user_management,
258+
client_id=TEST_CONSTANTS["CLIENT_ID"],
259+
session_data=session_data,
260+
cookie_password=TEST_CONSTANTS["COOKIE_PASSWORD"]
261+
)
262+
275263
with (
276-
patch(
277-
'workos.session.PyJWKClient',
278-
return_value=mock_jwks
264+
patch.object(
265+
session,
266+
"is_valid_jwt",
267+
return_value=True
279268
),
280-
):
281-
session = SessionModule(
282-
user_management=mock_user_management,
283-
client_id=TEST_CONSTANTS["CLIENT_ID"],
284-
session_data=session_data,
285-
cookie_password=TEST_CONSTANTS["COOKIE_PASSWORD"]
286-
)
287-
288-
with (
289-
patch.object(
290-
session,
291-
"is_valid_jwt",
292-
return_value=True
293-
),
294-
patch(
295-
"jwt.decode",
296-
return_value={
297-
"sid": TEST_CONSTANTS["SESSION_ID"],
298-
"org_id": TEST_CONSTANTS["ORGANIZATION_ID"],
299-
"role": "admin",
300-
"permissions": ["read"],
301-
"entitlements": ["feature_1"]
302-
}
303-
)
304-
):
305-
response = session.refresh()
306-
307-
assert isinstance(response, RefreshWithSessionCookieSuccessResponse)
308-
assert response.authenticated is True
309-
assert response.user.id == test_user["id"]
310-
311-
# Verify the refresh token was used correctly
312-
mock_user_management.authenticate_with_refresh_token.assert_called_once_with(
313-
refresh_token="refresh_token_12345",
314-
organization_id=None,
315-
session={
316-
"seal_session": True,
317-
"cookie_password": TEST_CONSTANTS["COOKIE_PASSWORD"]
269+
patch(
270+
"jwt.decode",
271+
return_value={
272+
"sid": TEST_CONSTANTS["SESSION_ID"],
273+
"org_id": TEST_CONSTANTS["ORGANIZATION_ID"],
274+
"role": "admin",
275+
"permissions": ["read"],
276+
"entitlements": ["feature_1"]
318277
}
319278
)
279+
):
280+
response = session.refresh()
281+
282+
assert isinstance(response, RefreshWithSessionCookieSuccessResponse)
283+
assert response.authenticated is True
284+
assert response.user.id == test_user["id"]
285+
286+
# Verify the refresh token was used correctly
287+
mock_user_management.authenticate_with_refresh_token.assert_called_once_with(
288+
refresh_token="refresh_token_12345",
289+
organization_id=None,
290+
session={
291+
"seal_session": True,
292+
"cookie_password": TEST_CONSTANTS["COOKIE_PASSWORD"]
293+
}
294+
)
320295

321296

322297
def test_seal_data(TEST_CONSTANTS):

0 commit comments

Comments
 (0)