import unittest import httpx from pwo import async_test from kaya.core import HttpContext, KayaApp from kaya.session import InMemorySessionStore, Session, SessionMixin class FakeClock: def __init__(self, start: float = 0.0) -> None: self._now = start def __call__(self) -> float: return self._now def advance(self, seconds: float) -> None: self._now += seconds class SessionIndexUnitTest(unittest.TestCase): def test_add_and_remove_index(self) -> None: session = Session() session.add_index('sub', 'alice') self.assertEqual(frozenset({('sub', 'alice')}), session.indexes) self.assertEqual(frozenset({('sub', 'alice')}), session.added_indexes) self.assertTrue(session.modified) session.remove_index('sub', 'alice') self.assertEqual(frozenset(), session.indexes) self.assertEqual(frozenset(), session.added_indexes) self.assertEqual(frozenset(), session.removed_indexes) def test_remove_persisted_index(self) -> None: session = Session(indexes={('sub', 'alice')}) self.assertFalse(session.modified) session.remove_index('sub', 'alice') self.assertEqual(frozenset(), session.indexes) self.assertEqual(frozenset({('sub', 'alice')}), session.removed_indexes) session.add_index('sub', 'alice') self.assertEqual(frozenset({('sub', 'alice')}), session.indexes) self.assertEqual(frozenset(), session.removed_indexes) def test_payload_round_trip(self) -> None: session = Session('abc', {'user': 'alice'}, indexes={('sub', 'alice')}) loaded = Session.from_payload('abc', session.to_payload()) self.assertEqual({'user': 'alice'}, dict(loaded)) self.assertEqual(frozenset({('sub', 'alice')}), loaded.indexes) def test_legacy_payload(self) -> None: loaded = Session.from_payload('abc', {'user': 'alice'}) self.assertEqual({'user': 'alice'}, dict(loaded)) self.assertEqual(frozenset(), loaded.indexes) class InMemorySessionIndexTest(unittest.TestCase): @async_test async def test_save_load_and_find(self) -> None: store = InMemorySessionStore() self.assertTrue(store.supports_indexes) session = Session('abc', {'user': 'alice'}) session.add_index('sub', 'alice') await store.save('abc', session) session.commit_indexes() loaded = await store.load('abc') assert loaded is not None self.assertEqual(frozenset({('sub', 'alice')}), loaded.indexes) self.assertEqual(['abc'], await store.find_sessions('sub', 'alice')) self.assertEqual([], await store.find_sessions('sub', 'bob')) @async_test async def test_delete_cleans_reverse_index(self) -> None: store = InMemorySessionStore() session = Session('abc', indexes={('sub', 'alice')}) await store.save('abc', session) await store.delete('abc') self.assertEqual([], await store.find_sessions('sub', 'alice')) self.assertEqual({}, store._session_indexes) @async_test async def test_delete_by_index_deletes_all_sessions(self) -> None: store = InMemorySessionStore() await store.save('a', Session('a', indexes={('sub', 'alice')})) await store.save('b', Session('b', indexes={('sub', 'alice')})) await store.save('c', Session('c', indexes={('sub', 'bob')})) self.assertEqual(2, await store.delete_by_index('sub', 'alice')) self.assertIsNone(await store.load('a')) self.assertIsNone(await store.load('b')) self.assertIsNotNone(await store.load('c')) @async_test async def test_expiry_prunes_index(self) -> None: clock = FakeClock() store = InMemorySessionStore(clock=clock) await store.save( 'abc', Session('abc', indexes={('sub', 'alice')}), max_age=60, ) clock.advance(61) self.assertEqual([], await store.find_sessions('sub', 'alice')) self.assertEqual({}, store._session_indexes) @async_test async def test_removed_index_is_not_readded(self) -> None: store = InMemorySessionStore() session = Session('abc', indexes={('sub', 'alice')}) await store.save('abc', session) session.commit_indexes() session.remove_index('sub', 'alice') await store.save('abc', session) self.assertEqual([], await store.find_sessions('sub', 'alice')) self.assertEqual(frozenset(), session.indexes) class SessionIndexIntegrationTest(unittest.TestCase): def setUp(self) -> None: self.store = InMemorySessionStore() self.mixin = SessionMixin(self.store) self.app = KayaApp(mixins=[self.mixin]) @self.app.GET('/login') async def login(ctx: HttpContext) -> None: ctx.session['user'] = 'alice' ctx.session.add_index('sub', 'alice') await ctx.send_str(200, 'ok') @self.app.GET('/rotate') async def rotate(ctx: HttpContext) -> None: ctx.session.regenerate_id() await ctx.send_str(200, 'ok') @self.app.GET('/logout') async def logout(ctx: HttpContext) -> None: ctx.session.invalidate() await ctx.send_str(200, 'ok') @async_test async def test_rotation_preserves_index(self) -> None: self.assertIs(self.store, self.mixin.store) transport = httpx.ASGITransport(app=self.app) async with httpx.AsyncClient(transport=transport, base_url='http://127.0.0.1:80') as client: await client.get('/login') old_id = client.cookies['session_id'] self.assertEqual([old_id], await self.store.find_sessions('sub', 'alice')) await client.get('/rotate') new_id = client.cookies['session_id'] self.assertNotEqual(old_id, new_id) self.assertEqual([new_id], await self.store.find_sessions('sub', 'alice')) @async_test async def test_invalidation_cleans_index(self) -> None: transport = httpx.ASGITransport(app=self.app) async with httpx.AsyncClient(transport=transport, base_url='http://127.0.0.1:80') as client: await client.get('/login') await client.get('/logout') self.assertEqual([], await self.store.find_sessions('sub', 'alice')) if __name__ == '__main__': unittest.main()