123 lines
4.9 KiB
Python
123 lines
4.9 KiB
Python
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()
|