Skip to content
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 7 additions & 0 deletions tests/test_session.py
Original file line number Diff line number Diff line change
Expand Up @@ -51,6 +51,7 @@ def session_constants(self):
"roles": ["admin"],
"permissions": ["read"],
"entitlements": ["feature_1"],
"feature_flags": ["flag1", "flag2"],
"exp": int(current_datetime.timestamp()) + 3600,
"iat": int(current_datetime.timestamp()),
}
Expand Down Expand Up @@ -244,6 +245,7 @@ def test_authenticate_success(self, session_constants, mock_user_management):
"roles": ["admin"],
"permissions": ["read"],
"entitlements": ["feature_1"],
"feature_flags": ["flag1", "flag2"],
}

with patch.object(Session, "unseal_data", return_value=mock_session), patch(
Expand All @@ -263,6 +265,7 @@ def test_authenticate_success(self, session_constants, mock_user_management):
assert response.roles == ["admin"]
assert response.permissions == ["read"]
assert response.entitlements == ["feature_1"]
assert response.feature_flags == ["flag1", "flag2"]
assert response.user.id == session_constants["USER_ID"]
assert response.impersonator is None

Expand Down Expand Up @@ -312,6 +315,7 @@ def test_authenticate_success_with_roles(
"roles": ["admin", "member"],
"permissions": ["read", "write"],
"entitlements": ["feature_1"],
"feature_flags": ["flag1", "flag2"],
}

with patch.object(Session, "unseal_data", return_value=mock_session), patch(
Expand All @@ -331,6 +335,7 @@ def test_authenticate_success_with_roles(
assert response.roles == ["admin", "member"]
assert response.permissions == ["read", "write"]
assert response.entitlements == ["feature_1"]
assert response.feature_flags == ["flag1", "flag2"]
assert response.user.id == session_constants["USER_ID"]
assert response.impersonator is None

Expand Down Expand Up @@ -410,6 +415,7 @@ def test_refresh_success(self, session_constants, mock_user_management):
"roles": ["admin"],
"permissions": ["read"],
"entitlements": ["feature_1"],
"feature_flags": ["flag1", "flag2"],
},
):
response = session.refresh()
Expand Down Expand Up @@ -511,6 +517,7 @@ async def test_refresh_success(self, session_constants, mock_user_management):
"roles": ["admin"],
"permissions": ["read"],
"entitlements": ["feature_1"],
"feature_flags": ["flag1", "flag2"],
},
):
response = await session.refresh()
Expand Down
4 changes: 4 additions & 0 deletions workos/session.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
from functools import lru_cache
import json
from typing import Any, Dict, Optional, Union, cast

import jwt
from jwt import PyJWKClient
from cryptography.fernet import Fernet
Expand Down Expand Up @@ -107,6 +108,7 @@ def authenticate(
entitlements=decoded.get("entitlements", None),
user=session["user"],
impersonator=session.get("impersonator", None),
feature_flags=decoded.get("feature_flags", None),
)

def refresh(
Expand Down Expand Up @@ -235,6 +237,7 @@ def refresh(
entitlements=decoded.get("entitlements", None),
user=auth_response.user,
impersonator=auth_response.impersonator,
feature_flags=decoded.get("feature_flags", None),
)
except Exception as e:
return RefreshWithSessionCookieErrorResponse(
Expand Down Expand Up @@ -326,6 +329,7 @@ async def refresh(
entitlements=decoded.get("entitlements", None),
user=auth_response.user,
impersonator=auth_response.impersonator,
feature_flags=decoded.get("feature_flags", None),
)
except Exception as e:
return RefreshWithSessionCookieErrorResponse(
Expand Down
1 change: 1 addition & 0 deletions workos/types/user_management/session.py
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@ class AuthenticateWithSessionCookieSuccessResponse(WorkOSModel):
user: User
impersonator: Optional[Impersonator] = None
entitlements: Optional[Sequence[str]] = None
feature_flags: Optional[Sequence[str]] = None


class AuthenticateWithSessionCookieErrorResponse(WorkOSModel):
Expand Down