Add OIDC back-channel logout with session indexes

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

@@ -0,0 +1,122 @@
+import pickle
+import unittest
+
+from pwo import async_test
+from test_store import FakeClock, StubMemcacheClient
+
+from kaya.session import Session
+from kaya.session.memcache import MemcacheSessionStore
+
+
+class MemcacheSessionIndexTest(unittest.TestCase):
+
+    def setUp(self) -> None:
+        self.clock = FakeClock()
+        self.client = StubMemcacheClient(self.clock)
+        self.store = MemcacheSessionStore(self.client, clock=self.clock)  # type: ignore[arg-type]
+
+    @async_test
+    async def test_save_load_and_find(self) -> None:
+        self.assertTrue(self.store.supports_indexes)
+        await self.store.save('abc', Session('abc', {'user': 'alice'}, indexes={('sub', 'alice')}))
+        loaded = await self.store.load('abc')
+        assert loaded is not None
+        self.assertEqual({'user': 'alice'}, dict(loaded))
+        self.assertEqual(frozenset({('sub', 'alice')}), loaded.indexes)
+        self.assertEqual(['abc'], await self.store.find_sessions('sub', 'alice'))
+
+    @async_test
+    async def test_rotation_preserves_index(self) -> None:
+        await self.store.save('old', Session('old', indexes={('sub', 'alice')}))
+        loaded = await self.store.load('old')
+        assert loaded is not None
+        loaded.set_id('new')
+        await self.store.delete('old')
+        await self.store.save('new', loaded)
+        self.assertEqual(['new'], await self.store.find_sessions('sub', 'alice'))
+
+    @async_test
+    async def test_delete_cleans_reverse_index(self) -> None:
+        await self.store.save('abc', Session('abc', indexes={('sub', 'alice')}))
+        await self.store.delete('abc')
+        self.assertEqual([], await self.store.find_sessions('sub', 'alice'))
+
+    @async_test
+    async def test_delete_by_index_deletes_all_sessions(self) -> None:
+        await self.store.save('a', Session('a', indexes={('sub', 'alice')}))
+        await self.store.save('b', Session('b', indexes={('sub', 'alice')}))
+        await self.store.save('c', Session('c', indexes={('sub', 'bob')}))
+        self.assertEqual(2, await self.store.delete_by_index('sub', 'alice'))
+        self.assertIsNone(await self.store.load('a'))
+        self.assertIsNone(await self.store.load('b'))
+        self.assertIsNotNone(await self.store.load('c'))
+
+    @async_test
+    async def test_stale_member_is_pruned(self) -> None:
+        await self.store.save('abc', Session('abc', indexes={('sub', 'alice')}))
+        await self.client.delete(b'kaya:session:abc')
+        self.assertEqual([], await self.store.find_sessions('sub', 'alice'))
+        index_key = self.store._index_key('sub', 'alice')
+        payload = await self.client.get(index_key)
+        assert payload is not None
+        members, _ = pickle.loads(payload)
+        self.assertEqual(set(), members)
+
+    @async_test
+    async def test_index_ttl_tracks_session_max_age(self) -> None:
+        await self.store.save(
+            'abc',
+            Session('abc', indexes={('sub', 'alice')}),
+            max_age=60,
+        )
+        index_key = self.store._index_key('sub', 'alice')
+        self.assertIn(index_key, self.client._expires)
+        self.clock.advance(61)
+        await self.store.save(
+            'abc',
+            Session('abc', indexes={('sub', 'alice')}),
+            max_age=60,
+        )
+        self.assertIsNotNone(await self.client.get(index_key))
+
+    @async_test
+    async def test_legacy_payload_is_upgraded_on_save(self) -> None:
+        await self.client.set(b'kaya:session:abc', pickle.dumps({'foo': 'bar'}))
+        loaded = await self.store.load('abc')
+        assert loaded is not None
+        self.assertEqual({'foo': 'bar'}, dict(loaded))
+        self.assertEqual(frozenset(), loaded.indexes)
+
+        loaded.add_index('sub', 'alice')
+        await self.store.save('abc', loaded)
+        raw = await self.client.get(b'kaya:session:abc')
+        assert raw is not None
+        payload = pickle.loads(raw)
+        self.assertEqual(1, payload['__kaya_session_v1__'])
+        self.assertEqual({'foo': 'bar'}, payload['data'])
+        self.assertEqual(['abc'], await self.store.find_sessions('sub', 'alice'))
+
+    @async_test
+    async def test_cas_conflict_is_retried(self) -> None:
+        await self.store.save('existing', Session('existing', indexes={('sub', 'alice')}))
+        original_cas = self.client.cas
+        attempts = 0
+
+        async def fail_once(key: bytes, value: bytes, cas_token: int, exptime: int = 0) -> bool:
+            nonlocal attempts
+            attempts += 1
+            if attempts == 1:
+                return False
+            return await original_cas(key, value, cas_token, exptime)
+
+        self.client.cas = fail_once  # type: ignore[method-assign]
+        await self.store.save('abc', Session('abc', indexes={('sub', 'alice')}))
+        self.assertEqual(
+            {'existing', 'abc'},
+            set(await self.store.find_sessions('sub', 'alice')),
+        )
+        self.assertGreaterEqual(attempts, 2)
+
+
+if __name__ == '__main__':
+    unittest.main()