Make tests and mypy happy

Author: PaulAsjes

tests/test_session.py

@@ -5,7 +5,7 @@
 from datetime import datetime, timezone
 
 from tests.conftest import with_jwks_mock
-from workos.session import SessionModule
+from workos.session import Session
 from workos.types.user_management.authentication_response import (
     RefreshTokenAuthenticationResponse,
 )
@@ -73,7 +73,7 @@ def mock_user_management():
 
 @with_jwks_mock
 def test_initialize_session_module(TEST_CONSTANTS, mock_user_management):
-    session = SessionModule(
+    session = Session(
         user_management=mock_user_management,
         client_id=TEST_CONSTANTS["CLIENT_ID"],
         session_data=TEST_CONSTANTS["SESSION_DATA"],
@@ -87,7 +87,7 @@ def test_initialize_session_module(TEST_CONSTANTS, mock_user_management):
 @with_jwks_mock
 def test_initialize_without_cookie_password(TEST_CONSTANTS, mock_user_management):
     with pytest.raises(ValueError, match="cookie_password is required"):
-        SessionModule(
+        Session(
             user_management=mock_user_management,
             client_id=TEST_CONSTANTS["CLIENT_ID"],
             session_data=TEST_CONSTANTS["SESSION_DATA"],
@@ -97,7 +97,7 @@ def test_initialize_without_cookie_password(TEST_CONSTANTS, mock_user_management
 
 @with_jwks_mock
 def test_authenticate_no_session_cookie_provided(TEST_CONSTANTS, mock_user_management):
-    session = SessionModule(
+    session = Session(
         user_management=mock_user_management,
         client_id=TEST_CONSTANTS["CLIENT_ID"],
         session_data=None,
@@ -114,7 +114,7 @@ def test_authenticate_no_session_cookie_provided(TEST_CONSTANTS, mock_user_manag
 
 @with_jwks_mock
 def test_authenticate_invalid_session_cookie(TEST_CONSTANTS, mock_user_management):
-    session = SessionModule(
+    session = Session(
         user_management=mock_user_management,
         client_id=TEST_CONSTANTS["CLIENT_ID"],
         session_data="invalid_session_data",
@@ -131,10 +131,10 @@ def test_authenticate_invalid_session_cookie(TEST_CONSTANTS, mock_user_managemen
 
 @with_jwks_mock
 def test_authenticate_invalid_jwt(TEST_CONSTANTS, mock_user_management):
-    invalid_session_data = SessionModule.seal_data(
+    invalid_session_data = Session.seal_data(
         {"access_token": "invalid_session_data"}, TEST_CONSTANTS["COOKIE_PASSWORD"]
     )
-    session = SessionModule(
+    session = Session(
         user_management=mock_user_management,
         client_id=TEST_CONSTANTS["CLIENT_ID"],
         session_data=invalid_session_data,
@@ -143,12 +143,14 @@ def test_authenticate_invalid_jwt(TEST_CONSTANTS, mock_user_management):
 
     response = session.authenticate()
 
+    print(response)
+
     assert response.reason == AuthenticateWithSessionCookieFailureReason.INVALID_JWT
 
 
 @with_jwks_mock
 def test_authenticate_success(TEST_CONSTANTS, mock_user_management):
-    session = SessionModule(
+    session = Session(
         user_management=mock_user_management,
         client_id=TEST_CONSTANTS["CLIENT_ID"],
         session_data=TEST_CONSTANTS["SESSION_DATA"],
@@ -192,7 +194,7 @@ def test_authenticate_success(TEST_CONSTANTS, mock_user_management):
 
     with (
         # Mock unsealing the session data
-        patch.object(SessionModule, "unseal_data", return_value=mock_session),
+        patch.object(Session, "unseal_data", return_value=mock_session),
         # Mock JWT validation
         patch.object(session, "is_valid_jwt", return_value=True),
         # Mock JWT decoding
@@ -219,7 +221,7 @@ def test_authenticate_success(TEST_CONSTANTS, mock_user_management):
 
 @with_jwks_mock
 def test_refresh_invalid_session_cookie(TEST_CONSTANTS, mock_user_management):
-    session = SessionModule(
+    session = Session(
         user_management=mock_user_management,
         client_id=TEST_CONSTANTS["CLIENT_ID"],
         session_data="invalid_session_data",
@@ -248,7 +250,7 @@ def test_refresh_success(TEST_CONSTANTS, mock_user_management):
         "updated_at": TEST_CONSTANTS["CURRENT_TIMESTAMP"],
     }
 
-    session_data = SessionModule.seal_data(
+    session_data = Session.seal_data(
         {"refresh_token": "refresh_token_12345", "user": test_user},
         TEST_CONSTANTS["COOKIE_PASSWORD"],
     )
@@ -264,7 +266,7 @@ def test_refresh_success(TEST_CONSTANTS, mock_user_management):
         RefreshTokenAuthenticationResponse(**mock_response)
     )
 
-    session = SessionModule(
+    session = Session(
         user_management=mock_user_management,
         client_id=TEST_CONSTANTS["CLIENT_ID"],
         session_data=session_data,
@@ -303,16 +305,19 @@ def test_refresh_success(TEST_CONSTANTS, mock_user_management):
 
 def test_seal_data(TEST_CONSTANTS):
     test_data = {"test": "data"}
-    sealed = SessionModule.seal_data(test_data, TEST_CONSTANTS["COOKIE_PASSWORD"])
+    sealed = Session.seal_data(test_data, TEST_CONSTANTS["COOKIE_PASSWORD"])
     assert isinstance(sealed, str)
 
     # Test unsealing
-    unsealed = SessionModule.unseal_data(sealed, TEST_CONSTANTS["COOKIE_PASSWORD"])
+    unsealed = Session.unseal_data(sealed, TEST_CONSTANTS["COOKIE_PASSWORD"])
+    print(test_data)
+    print(unsealed)
+
     assert unsealed == test_data
 
 
 def test_unseal_invalid_data(TEST_CONSTANTS):
     with pytest.raises(Exception):  # Adjust exception type based on your implementation
-        SessionModule.unseal_data(
+        Session.unseal_data(
             "invalid_sealed_data", TEST_CONSTANTS["COOKIE_PASSWORD"]
         )

workos/session.py

@@ -1,5 +1,6 @@
+from __future__ import annotations
 import json
-from typing import Any, Dict, List, Optional, Union
+from typing import Any, Dict, Optional, Union, cast
 import jwt
 from jwt import PyJWKClient
 from cryptography.fernet import Fernet
@@ -12,12 +13,11 @@
     RefreshWithSessionCookieSuccessResponse,
 )
 
-
-class SessionModule:
+class Session:
     def __init__(
         self,
         *,
-        user_management: Any,
+        user_management: "UserManagementModule", # type: ignore
         client_id: str,
         session_data: str,
         cookie_password: str,
@@ -42,7 +42,7 @@ def authenticate(
         AuthenticateWithSessionCookieSuccessResponse,
         AuthenticateWithSessionCookieErrorResponse,
     ]:
-        if self.session_data is None:
+        if self.session_data is None or self.session_data == "":
             return AuthenticateWithSessionCookieErrorResponse(
                 authenticated=False,
                 reason=AuthenticateWithSessionCookieFailureReason.NO_SESSION_COOKIE_PROVIDED,
@@ -56,7 +56,7 @@ def authenticate(
                 reason=AuthenticateWithSessionCookieFailureReason.INVALID_SESSION_COOKIE,
             )
 
-        if not session["access_token"]:
+        if not session.get("access_token", None):
             return AuthenticateWithSessionCookieErrorResponse(
                 authenticated=False,
                 reason=AuthenticateWithSessionCookieFailureReason.INVALID_SESSION_COOKIE,
@@ -84,14 +84,15 @@ def authenticate(
             impersonator=session.get("impersonator", None),
         )
 
-    def refresh(self, options: Optional[Dict[str, Any]] = None) -> Union[
+    def refresh(
+        self, *, organization_id: Optional[str] = None, cookie_password: Optional[str] = None
+    ) -> Union[
         RefreshWithSessionCookieSuccessResponse,
         RefreshWithSessionCookieErrorResponse,
     ]:
         cookie_password = (
-            self.cookie_password if options is None else options.get("cookie_password")
+            self.cookie_password if cookie_password is None else cookie_password
         )
-        organization_id = None if options is None else options.get("organization_id")
 
         try:
             session = self.unseal_data(self.session_data, cookie_password)
@@ -101,7 +102,7 @@ def refresh(self, options: Optional[Dict[str, Any]] = None) -> Union[
                 reason=AuthenticateWithSessionCookieFailureReason.INVALID_SESSION_COOKIE,
             )
 
-        if not session["refresh_token"] or not session["user"]:
+        if not session.get("refresh_token", None) or not session.get("user", None):
             return RefreshWithSessionCookieErrorResponse(
                 authenticated=False,
                 reason=AuthenticateWithSessionCookieFailureReason.INVALID_SESSION_COOKIE,
@@ -144,12 +145,14 @@ def refresh(self, options: Optional[Dict[str, Any]] = None) -> Union[
     def get_logout_url(self) -> str:
         auth_response = self.authenticate()
 
-        if not auth_response.authenticated:
+        if isinstance(auth_response, AuthenticateWithSessionCookieErrorResponse):
             raise ValueError(
                 f"Failed to extract session ID for logout URL: {auth_response.reason}"
             )
 
-        return self.user_management.get_logout_url(session_id=auth_response.session_id)
+        result = self.user_management.get_logout_url(session_id=auth_response.session_id)
+        return str(result)
+
 
     def is_valid_jwt(self, token: str) -> bool:
         try:
@@ -171,4 +174,5 @@ def unseal_data(sealed_data: str, key: str) -> Dict[str, Any]:
         fernet = Fernet(key)
         # Convert string back to bytes before decryption
         encrypted_bytes = sealed_data.encode("utf-8")
-        return json.loads(fernet.decrypt(encrypted_bytes).decode())
+        decrypted_str = fernet.decrypt(encrypted_bytes).decode()
+        return cast(Dict[str, Any], json.loads(decrypted_str))

workos/types/user_management/session.py

@@ -1,6 +1,6 @@
 from typing import List, Optional, TypedDict, Union
 from enum import Enum
-
+from typing_extensions import Literal
 from workos.types.user_management.impersonator import Impersonator
 from workos.types.user_management.user import User
 from workos.types.workos_model import WorkOSModel
@@ -13,7 +13,7 @@ class AuthenticateWithSessionCookieFailureReason(Enum):
 
 
 class AuthenticateWithSessionCookieSuccessResponse(WorkOSModel):
-    authenticated: bool = True
+    authenticated: Literal[True]
     session_id: str
     organization_id: Optional[str] = None
     role: Optional[str] = None
@@ -24,7 +24,7 @@ class AuthenticateWithSessionCookieSuccessResponse(WorkOSModel):
 
 
 class AuthenticateWithSessionCookieErrorResponse(WorkOSModel):
-    authenticated: bool = False
+    authenticated: Literal[False]
     reason: Union[AuthenticateWithSessionCookieFailureReason, str]
 
 
@@ -35,10 +35,10 @@ class RefreshWithSessionCookieSuccessResponse(
 
 
 class RefreshWithSessionCookieErrorResponse(WorkOSModel):
-    authenticated: bool = False
+    authenticated: Literal[False]
     reason: Union[AuthenticateWithSessionCookieFailureReason, str]
 
 
-class SessionConfig(TypedDict):
+class SessionConfig(TypedDict, total=False):
     seal_session: bool
     cookie_password: str