diff --git a/sdks/implementations/python/tests/test_token_store.py b/sdks/implementations/python/tests/test_token_store.py new file mode 100644 index 000000000..363459b8b --- /dev/null +++ b/sdks/implementations/python/tests/test_token_store.py @@ -0,0 +1,406 @@ +"""Tests for the token store subsystem: ABC, concrete stores, registry, and CAS algorithm.""" + +from __future__ import annotations + +import asyncio +import base64 +import json +import time +from typing import Mapping +from unittest.mock import MagicMock, patch + +import pytest + +from stack_auth._token_store import ( + ExplicitTokenStore, + MemoryTokenStore, + RequestTokenStore, + TokenStore, + _decode_jwt_payload, + _is_expired, + _is_fresh_enough, + get_or_fetch_likely_valid_tokens_async, + get_or_fetch_likely_valid_tokens_sync, + resolve_token_store, +) + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +def _make_jwt(claims: dict) -> str: + """Build a fake JWT with the given payload claims (no real signature).""" + header = base64.urlsafe_b64encode( + json.dumps({"typ": "JWT", "alg": "RS256"}).encode() + ).rstrip(b"=").decode() + payload = base64.urlsafe_b64encode( + json.dumps(claims).encode() + ).rstrip(b"=").decode() + return f"{header}.{payload}.fake_signature" + + +def _make_fresh_token(now: float) -> str: + """Token that is fresh: exp > now+20 and iat > now-75.""" + return _make_jwt({"sub": "user_1", "exp": now + 60, "iat": now - 10}) + + +def _make_expiring_soon_token(now: float) -> str: + """Token that expires within 20s.""" + return _make_jwt({"sub": "user_1", "exp": now + 10, "iat": now - 10}) + + +def _make_old_token(now: float) -> str: + """Token issued > 75s ago (stale iat).""" + return _make_jwt({"sub": "user_1", "exp": now + 60, "iat": now - 100}) + + +def _make_expired_token(now: float) -> str: + """Token that is already expired.""" + return _make_jwt({"sub": "user_1", "exp": now - 10, "iat": now - 100}) + + +class _FakeRequest: + """Simple request-like object implementing the RequestLike protocol.""" + + def __init__(self, headers: dict[str, str]) -> None: + self._headers = headers + + @property + def headers(self) -> Mapping[str, str]: + return self._headers + + +# --------------------------------------------------------------------------- +# TokenStore ABC +# --------------------------------------------------------------------------- + +class TestTokenStoreABC: + def test_cannot_instantiate_abstract(self) -> None: + with pytest.raises(TypeError): + TokenStore() # type: ignore[abstract] + + def test_has_sync_and_async_locks(self) -> None: + store = MemoryTokenStore() + assert hasattr(store, "_sync_lock") + assert hasattr(store, "_async_lock") + + +# --------------------------------------------------------------------------- +# MemoryTokenStore +# --------------------------------------------------------------------------- + +class TestMemoryTokenStore: + def test_stores_and_retrieves_tokens(self) -> None: + store = MemoryTokenStore() + assert store.get_stored_access_token() is None + assert store.get_stored_refresh_token() is None + + store.compare_and_set(None, "rt_1", "at_1") # type: ignore[arg-type] + # compare against None should match initial state + assert store.get_stored_access_token() == "at_1" + assert store.get_stored_refresh_token() == "rt_1" + + def test_compare_and_set_updates_when_matching(self) -> None: + store = MemoryTokenStore() + store.compare_and_set(None, "rt_1", "at_1") # type: ignore[arg-type] + store.compare_and_set("rt_1", "rt_2", "at_2") + assert store.get_stored_access_token() == "at_2" + assert store.get_stored_refresh_token() == "rt_2" + + def test_compare_and_set_does_not_update_when_mismatched(self) -> None: + store = MemoryTokenStore() + store.compare_and_set(None, "rt_1", "at_1") # type: ignore[arg-type] + store.compare_and_set("wrong_rt", "rt_2", "at_2") + assert store.get_stored_access_token() == "at_1" + assert store.get_stored_refresh_token() == "rt_1" + + +# --------------------------------------------------------------------------- +# ExplicitTokenStore +# --------------------------------------------------------------------------- + +class TestExplicitTokenStore: + def test_initializes_from_tokens(self) -> None: + store = ExplicitTokenStore(access_token="at_1", refresh_token="rt_1") + assert store.get_stored_access_token() == "at_1" + assert store.get_stored_refresh_token() == "rt_1" + + def test_supports_cas_update(self) -> None: + store = ExplicitTokenStore(access_token="at_1", refresh_token="rt_1") + store.compare_and_set("rt_1", "rt_2", "at_2") + assert store.get_stored_access_token() == "at_2" + assert store.get_stored_refresh_token() == "rt_2" + + def test_cas_does_not_update_when_mismatched(self) -> None: + store = ExplicitTokenStore(access_token="at_1", refresh_token="rt_1") + store.compare_and_set("wrong", "rt_2", "at_2") + assert store.get_stored_access_token() == "at_1" + assert store.get_stored_refresh_token() == "rt_1" + + +# --------------------------------------------------------------------------- +# RequestTokenStore +# --------------------------------------------------------------------------- + +class TestRequestTokenStore: + def test_extracts_tokens_from_header(self) -> None: + header_value = json.dumps({"accessToken": "at_req", "refreshToken": "rt_req"}) + request = _FakeRequest({"x-stack-auth": header_value}) + store = RequestTokenStore(request) + assert store.get_stored_access_token() == "at_req" + assert store.get_stored_refresh_token() == "rt_req" + + def test_returns_none_when_header_missing(self) -> None: + request = _FakeRequest({}) + store = RequestTokenStore(request) + assert store.get_stored_access_token() is None + assert store.get_stored_refresh_token() is None + + def test_returns_none_when_header_is_malformed_json(self) -> None: + request = _FakeRequest({"x-stack-auth": "not-json"}) + store = RequestTokenStore(request) + assert store.get_stored_access_token() is None + assert store.get_stored_refresh_token() is None + + def test_supports_cas_update(self) -> None: + header_value = json.dumps({"accessToken": "at_req", "refreshToken": "rt_req"}) + request = _FakeRequest({"x-stack-auth": header_value}) + store = RequestTokenStore(request) + store.compare_and_set("rt_req", "rt_new", "at_new") + assert store.get_stored_access_token() == "at_new" + assert store.get_stored_refresh_token() == "rt_new" + + +# --------------------------------------------------------------------------- +# Registry +# --------------------------------------------------------------------------- + +class TestTokenStoreRegistry: + def test_same_project_returns_same_instance(self) -> None: + store1 = resolve_token_store("memory", "proj_1") + store2 = resolve_token_store("memory", "proj_1") + assert store1 is store2 + + def test_different_projects_return_different_instances(self) -> None: + store1 = resolve_token_store("memory", "proj_1") + store2 = resolve_token_store("memory", "proj_2") + assert store1 is not store2 + + def test_resolve_explicit_token_store(self) -> None: + store = resolve_token_store( + {"access_token": "at", "refresh_token": "rt"}, "proj" + ) + assert isinstance(store, ExplicitTokenStore) + assert store.get_stored_access_token() == "at" + assert store.get_stored_refresh_token() == "rt" + + def test_resolve_request_token_store(self) -> None: + header_value = json.dumps({"accessToken": "at", "refreshToken": "rt"}) + request = _FakeRequest({"x-stack-auth": header_value}) + store = resolve_token_store(request, "proj") + assert isinstance(store, RequestTokenStore) + + def test_resolve_none(self) -> None: + assert resolve_token_store(None, "proj") is None + + +# --------------------------------------------------------------------------- +# Helper functions: _is_fresh_enough, _is_expired, _decode_jwt_payload +# --------------------------------------------------------------------------- + +class TestIsFreshEnough: + def test_returns_true_for_fresh_token(self) -> None: + now = time.time() + token = _make_fresh_token(now) + with patch("stack_auth._token_store.time") as mock_time: + mock_time.time.return_value = now + assert _is_fresh_enough(token) is True + + def test_returns_false_for_expiring_soon(self) -> None: + now = time.time() + token = _make_expiring_soon_token(now) + with patch("stack_auth._token_store.time") as mock_time: + mock_time.time.return_value = now + assert _is_fresh_enough(token) is False + + def test_returns_false_for_old_iat(self) -> None: + now = time.time() + token = _make_old_token(now) + with patch("stack_auth._token_store.time") as mock_time: + mock_time.time.return_value = now + assert _is_fresh_enough(token) is False + + def test_returns_false_for_none(self) -> None: + assert _is_fresh_enough(None) is False + + +class TestDecodeJwtPayload: + def test_decodes_valid_jwt(self) -> None: + token = _make_jwt({"sub": "user_1", "exp": 9999999999}) + claims = _decode_jwt_payload(token) + assert claims is not None + assert claims["sub"] == "user_1" + + def test_returns_none_for_invalid(self) -> None: + assert _decode_jwt_payload("not-a-jwt") is None + + def test_returns_none_for_malformed_base64(self) -> None: + assert _decode_jwt_payload("a.!!!.c") is None + + +class TestIsExpired: + def test_expired_token(self) -> None: + now = time.time() + token = _make_expired_token(now) + with patch("stack_auth._token_store.time") as mock_time: + mock_time.time.return_value = now + assert _is_expired(token) is True + + def test_valid_token(self) -> None: + now = time.time() + token = _make_fresh_token(now) + with patch("stack_auth._token_store.time") as mock_time: + mock_time.time.return_value = now + assert _is_expired(token) is False + + +# --------------------------------------------------------------------------- +# CAS Refresh: Sync +# --------------------------------------------------------------------------- + +class TestCASRefreshSync: + def test_returns_fresh_access_token(self) -> None: + now = time.time() + store = MemoryTokenStore() + fresh_token = _make_fresh_token(now) + store.compare_and_set(None, "rt_1", fresh_token) # type: ignore[arg-type] + refresh_fn = MagicMock(return_value=(True, "new_at")) + + with patch("stack_auth._token_store.time") as mock_time: + mock_time.time.return_value = now + rt, at = get_or_fetch_likely_valid_tokens_sync(store, refresh_fn) + + assert rt == "rt_1" + assert at == fresh_token + refresh_fn.assert_not_called() + + def test_refreshes_when_expiring_soon(self) -> None: + now = time.time() + store = MemoryTokenStore() + expiring_token = _make_expiring_soon_token(now) + store.compare_and_set(None, "rt_1", expiring_token) # type: ignore[arg-type] + refresh_fn = MagicMock(return_value=(True, "new_at")) + + with patch("stack_auth._token_store.time") as mock_time: + mock_time.time.return_value = now + rt, at = get_or_fetch_likely_valid_tokens_sync(store, refresh_fn) + + assert rt == "rt_1" + assert at == "new_at" + refresh_fn.assert_called_once_with("rt_1") + + def test_refreshes_when_old_iat(self) -> None: + now = time.time() + store = MemoryTokenStore() + old_token = _make_old_token(now) + store.compare_and_set(None, "rt_1", old_token) # type: ignore[arg-type] + refresh_fn = MagicMock(return_value=(True, "new_at")) + + with patch("stack_auth._token_store.time") as mock_time: + mock_time.time.return_value = now + rt, at = get_or_fetch_likely_valid_tokens_sync(store, refresh_fn) + + assert rt == "rt_1" + assert at == "new_at" + refresh_fn.assert_called_once_with("rt_1") + + def test_clears_tokens_on_invalid_refresh(self) -> None: + now = time.time() + store = MemoryTokenStore() + expiring_token = _make_expiring_soon_token(now) + store.compare_and_set(None, "rt_1", expiring_token) # type: ignore[arg-type] + refresh_fn = MagicMock(return_value=(False, None)) + + with patch("stack_auth._token_store.time") as mock_time: + mock_time.time.return_value = now + rt, at = get_or_fetch_likely_valid_tokens_sync(store, refresh_fn) + + assert rt is None + assert at is None + # Store should be cleared + assert store.get_stored_refresh_token() is None + assert store.get_stored_access_token() is None + + def test_returns_access_without_refresh_when_not_expired(self) -> None: + now = time.time() + store = MemoryTokenStore() + fresh_token = _make_fresh_token(now) + # Set access token but no refresh token + store._access_token = fresh_token + refresh_fn = MagicMock() + + with patch("stack_auth._token_store.time") as mock_time: + mock_time.time.return_value = now + rt, at = get_or_fetch_likely_valid_tokens_sync(store, refresh_fn) + + assert rt is None + assert at == fresh_token + refresh_fn.assert_not_called() + + def test_returns_none_none_when_no_refresh_and_expired(self) -> None: + now = time.time() + store = MemoryTokenStore() + expired_token = _make_expired_token(now) + store._access_token = expired_token + refresh_fn = MagicMock() + + with patch("stack_auth._token_store.time") as mock_time: + mock_time.time.return_value = now + rt, at = get_or_fetch_likely_valid_tokens_sync(store, refresh_fn) + + assert rt is None + assert at is None + refresh_fn.assert_not_called() + + +# --------------------------------------------------------------------------- +# CAS Refresh: Async +# --------------------------------------------------------------------------- + +class TestCASRefreshAsync: + async def test_async_returns_fresh_access_token(self) -> None: + now = time.time() + store = MemoryTokenStore() + fresh_token = _make_fresh_token(now) + store.compare_and_set(None, "rt_1", fresh_token) # type: ignore[arg-type] + + async def refresh_fn(rt: str) -> tuple[bool, str | None]: + return (True, "new_at") + + with patch("stack_auth._token_store.time") as mock_time: + mock_time.time.return_value = now + rt, at = await get_or_fetch_likely_valid_tokens_async(store, refresh_fn) + + assert rt == "rt_1" + assert at == fresh_token + + async def test_async_refreshes_when_needed(self) -> None: + now = time.time() + store = MemoryTokenStore() + expiring_token = _make_expiring_soon_token(now) + store.compare_and_set(None, "rt_1", expiring_token) # type: ignore[arg-type] + called = False + + async def refresh_fn(rt: str) -> tuple[bool, str | None]: + nonlocal called + called = True + return (True, "new_at") + + with patch("stack_auth._token_store.time") as mock_time: + mock_time.time.return_value = now + rt, at = await get_or_fetch_likely_valid_tokens_async(store, refresh_fn) + + assert rt == "rt_1" + assert at == "new_at" + assert called