kaya/packages/kaya-session/tests/test_indexes.py

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()