test: Add more tests for session claim features

Author: KShivendu

tests/sessions/__init__.py

@@ -0,0 +1,75 @@
+import time
+from unittest.mock import MagicMock
+
+from supertokens_python import init
+from supertokens_python.framework import BaseRequest
+from supertokens_python.recipe import session
+from supertokens_python.recipe.session.asyncio import create_new_session
+from tests.utils import setup_function, teardown_function, start_st, min_api_version
+from .utils import (
+    st_init_args,
+    NoneClaim,
+    get_st_init_args,
+    session_functions_override_with_claim,
+    TrueClaim,
+)
+
+_ = setup_function  # type:ignore
+_ = teardown_function  # type:ignore
+
+timestamp = time.time()
+
+
+@min_api_version("2.13")
+async def test_create_access_token_payload_with_session_claims():
+    init(**st_init_args)  # type:ignore
+    start_st()
+
+    dummy_req: BaseRequest = MagicMock()
+    s = await create_new_session(dummy_req, "someId")
+
+    payload = s.get_access_token_payload()
+    assert payload["st-true"]["t"] > 0
+    payload["st-true"]["t"] = timestamp  # to avoid patching
+    assert payload == {"st-true": {"v": True, "t": timestamp}}
+
+
+@min_api_version("2.13")
+async def test_should_create_access_token_payload_with_session_claims_with_an_none_value():
+    init(**get_st_init_args(NoneClaim))  # type:ignore
+    start_st()
+
+    dummy_req: BaseRequest = MagicMock()
+    s = await create_new_session(dummy_req, "someId")
+
+    payload = s.get_access_token_payload()
+    assert payload == {}
+
+
+@min_api_version("2.13")
+async def test_should_merge_claims_and_passed_access_token_payload_obj():
+    new_st_init = {
+        **st_init_args,
+        "recipe_list": [
+            session.init(
+                override=session.InputOverrideConfig(
+                    functions=session_functions_override_with_claim(
+                        TrueClaim, {"user-custom-claim": "foo"}
+                    ),
+                )
+            ),
+        ],
+    }
+    init(**new_st_init)  # type:ignore
+    start_st()
+
+    dummy_req: BaseRequest = MagicMock()
+    s = await create_new_session(dummy_req, "someId")
+
+    payload = s.get_access_token_payload()
+    assert payload["st-true"]["t"] > 0
+    payload["st-true"]["t"] = timestamp  # to avoid patching
+    assert payload == {
+        "st-true": {"v": True, "t": timestamp},
+        "user-custom-claim": "foo",
+    }

tests/sessions/claims/test_create_new_session.py

@@ -0,0 +1,52 @@
+import time
+from unittest.mock import AsyncMock, patch
+
+from pytest import mark
+
+from supertokens_python.recipe.session.session_class import Session
+from tests.sessions.claims.utils import NoneClaim, TrueClaim
+
+timestamp = time.time()
+
+pytestmark = (
+    mark.asyncio
+)  # no need to apply @mark.asyncio on each test because of this!
+
+
+async def test_should_not_change_if_claim_fetch_value_returns_none():
+    recipe_implementation_mock = AsyncMock()
+    session = Session(
+        recipe_implementation_mock,
+        "test_access_token",
+        "test_session_handle",
+        "test_user_id",
+        {},
+    )
+    with patch.object(
+        Session,
+        "merge_into_access_token_payload",
+        wraps=session.merge_into_access_token_payload,
+    ) as mock:
+        await session.fetch_and_set_claim(NoneClaim)
+        mock.assert_called_once_with({}, None)
+
+
+async def test_should_update_if_claim_fetch_value_returns_value():
+    recipe_implementation_mock = AsyncMock()
+    session = Session(
+        recipe_implementation_mock,
+        "test_access_token",
+        "test_session_handle",
+        "test_user_id",
+        {},
+    )
+    with patch.object(
+        Session,
+        "merge_into_access_token_payload",
+        wraps=session.merge_into_access_token_payload,
+    ) as mock:
+        await session.fetch_and_set_claim(TrueClaim)
+        update, _ = mock.call_args_list[0].args
+        assert update["st-true"]["t"] > 0
+        update["st-true"]["t"] = timestamp
+        mock.assert_called_once_with({"st-true": {"t": timestamp, "v": True}}, None)

tests/sessions/claims/test_fetch_and_set_claim.py

@@ -0,0 +1,52 @@
+from unittest.mock import MagicMock
+
+from pytest import mark, raises
+
+from supertokens_python import init
+from supertokens_python.framework.request import BaseRequest
+from supertokens_python.recipe import session
+from supertokens_python.recipe.session.asyncio import (
+    create_new_session,
+    get_claim_value,
+)
+from supertokens_python.recipe.session.interfaces import SessionContainer
+from tests.utils import setup_function, teardown_function, start_st
+from .utils import TrueClaim, st_init_args
+
+_ = setup_function  # type:ignore
+_ = teardown_function  # type:ignore
+
+
+@mark.asyncio
+async def test_should_get_the_right_value():
+    init(**st_init_args)  # type:ignore
+    start_st()
+
+    dummy_req: BaseRequest = MagicMock()
+    s = await create_new_session(dummy_req, "someId")
+
+    res = await s.get_claim_value(TrueClaim)
+    assert res is True
+
+
+@mark.asyncio
+async def test_should_get_the_right_value_using_session_handle():
+    init(**st_init_args)  # type:ignore
+    start_st()
+
+    dummy_req: BaseRequest = MagicMock()
+    s: SessionContainer = await create_new_session(dummy_req, "someId")
+
+    res = await get_claim_value(s.get_handle(), TrueClaim)
+    assert res is True
+
+
+@mark.asyncio
+async def test_should_throw_for_non_existing_handle():
+    new_st_init = {**st_init_args, "recipe_list": [session.init()]}
+    init(**new_st_init)  # type: ignore
+    start_st()
+
+    with raises(Exception) as e:
+        _ = await get_claim_value("non_existing_handle", TrueClaim)
+        assert str(e) == "Session does not exist"

tests/sessions/claims/test_get_claims_value.py

@@ -0,0 +1,86 @@
+import time as real_time
+from unittest.mock import MagicMock, AsyncMock, patch
+
+from pytest import mark
+
+from supertokens_python import init
+from supertokens_python.framework import BaseRequest
+from supertokens_python.recipe.session import SessionContainer
+from supertokens_python.recipe.session.asyncio import create_new_session
+from supertokens_python.recipe.session.asyncio import (
+    remove_claim,
+    get_session_information,
+)
+from supertokens_python.recipe.session.session_class import Session
+from tests.sessions.claims.utils import TrueClaim
+from tests.utils import start_st, setup_function, teardown_function
+from .test_get_claims_value import st_init_args
+from .utils import time_patch_wrapper
+
+_ = setup_function  # type:ignore
+_ = teardown_function  # type:ignore
+
+timestamp = real_time.time()
+
+
+@mark.asyncio
+async def test_should_attempt_to_set_claim_to_none():
+    recipe_implementation_mock = AsyncMock()
+    session = Session(
+        recipe_implementation_mock,
+        "test_access_token",
+        "test_session_handle",
+        "test_user_id",
+        {},
+    )
+    with patch.object(
+        Session,
+        "merge_into_access_token_payload",
+        wraps=session.merge_into_access_token_payload,
+    ) as mock:
+        await session.remove_claim(TrueClaim)
+        mock.assert_called_once_with({"st-true": None}, None)
+
+
+@time_patch_wrapper
+async def test_should_clear_previously_set_claim(time_mock: MagicMock):
+    time_mock.time.return_value = timestamp  # type: ignore
+    init(**st_init_args)  # type:ignore
+    start_st()
+
+    dummy_req: BaseRequest = MagicMock()
+    s: SessionContainer = await create_new_session(dummy_req, "someId")
+
+    payload = s.get_access_token_payload()
+
+    assert payload == {"st-true": {"v": True, "t": timestamp}}
+
+
+@time_patch_wrapper
+async def test_should_clear_previously_set_claim_using_handle(time_mock: MagicMock):
+    time_mock.time.return_value = timestamp  # type: ignore
+    init(**st_init_args)  # type:ignore
+    start_st()
+
+    dummy_req: BaseRequest = MagicMock()
+    s: SessionContainer = await create_new_session(dummy_req, "someId")
+
+    payload = s.get_access_token_payload()
+    assert payload == {"st-true": {"v": True, "t": timestamp}}
+
+    res = await remove_claim(s.get_handle(), TrueClaim)
+    assert res is True
+
+    session_info = await get_session_information(s.get_handle())
+    assert session_info is not None
+    payload_after = session_info.access_token_payload
+    assert payload_after == {}
+
+
+@time_patch_wrapper
+async def test_should_work_ok_for_non_existing_handle(_time_mock: MagicMock):
+    init(**st_init_args)  # type:ignore
+    start_st()
+
+    res = await remove_claim("non-existing-handle", TrueClaim)
+    assert res is False

tests/sessions/claims/test_remove_claim.py

@@ -0,0 +1,102 @@
+import time
+from unittest.mock import AsyncMock, patch, MagicMock
+
+from pytest import mark
+
+from supertokens_python import init
+from supertokens_python.framework import BaseRequest
+from supertokens_python.recipe.session.asyncio import (
+    create_new_session,
+    get_session_information,
+    set_claim_value,
+)
+from supertokens_python.recipe.session.session_class import Session
+from tests.sessions.claims.utils import TrueClaim, st_init_args
+from tests.utils import setup_function, teardown_function
+from tests.utils import start_st
+
+_ = setup_function  # type:ignore
+_ = teardown_function  # type:ignore
+
+timestamp = time.time()
+
+pytestmark = (
+    mark.asyncio
+)  # no need to apply @mark.asyncio on each test because of this!
+
+
+async def test_should_merge_the_right_value():
+    recipe_implementation_mock = AsyncMock()
+    session = Session(
+        recipe_implementation_mock,
+        "test_access_token",
+        "test_session_handle",
+        "test_user_id",
+        {},
+    )
+    with patch.object(
+        Session,
+        "merge_into_access_token_payload",
+        wraps=session.merge_into_access_token_payload,
+    ) as mock:
+        await session.set_claim_value(TrueClaim, "NEW_TRUE")
+        update, _ = mock.call_args_list[0].args
+        assert update["st-true"]["t"] > 0
+        update["st-true"]["t"] = timestamp
+        mock.assert_called_once_with(
+            {"st-true": {"t": timestamp, "v": "NEW_TRUE"}}, None
+        )
+
+
+async def test_should_overwrite_claim_value():
+    init(**st_init_args)  # type: ignore
+    start_st()
+
+    dummy_req: BaseRequest = MagicMock()
+    s = await create_new_session(dummy_req, "someId")
+
+    payload = s.get_access_token_payload()
+    assert payload["st-true"]["t"] > 0
+    payload["st-true"]["t"] = timestamp
+    assert payload == {"st-true": {"t": timestamp, "v": True}}
+
+    await s.set_claim_value(TrueClaim, "NEW_TRUE")
+
+    # Payload should be updated now:
+    payload = s.get_access_token_payload()
+    assert payload["st-true"]["t"] > 0
+    payload["st-true"]["t"] = timestamp
+    assert payload == {"st-true": {"t": timestamp, "v": "NEW_TRUE"}}
+
+
+async def test_should_overwrite_claim_value_using_session_handle():
+    init(**st_init_args)  # type: ignore
+    start_st()
+
+    dummy_req: BaseRequest = MagicMock()
+    s = await create_new_session(dummy_req, "someId")
+
+    payload = s.get_access_token_payload()
+    assert payload["st-true"]["t"] > 0
+    payload["st-true"]["t"] = timestamp
+    assert payload == {"st-true": {"t": timestamp, "v": True}}
+
+    await set_claim_value(s.get_handle(), TrueClaim, "NEW_TRUE")
+
+    # Payload should be updated now:
+    # Note that the session var (s) still contains the old payload.
+    # We need to fetch the new one.
+    s = await get_session_information(s.get_handle())
+    assert s is not None
+    payload = s.access_token_payload
+    assert payload["st-true"]["t"] > 0
+    payload["st-true"]["t"] = timestamp
+    assert payload == {"st-true": {"t": timestamp, "v": "NEW_TRUE"}}
+
+
+async def test_should_work_ok_for_non_existing_handles():
+    init(**st_init_args)  # type: ignore
+    start_st()
+
+    res = await set_claim_value("non-existing-handle", TrueClaim, "NEW_TRUE")
+    assert res is False