From 1da2d67f105c35b96c9a22513bc44f00af5e2eb6 Mon Sep 17 00:00:00 2001 From: Smialy Date: Tue, 12 Sep 2017 18:46:54 +0200 Subject: [PATCH 1/2] Add key_factory attr for redis_storage --- aiohttp_session/redis_storage.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/aiohttp_session/redis_storage.py b/aiohttp_session/redis_storage.py index c3fe1209..9ffd55cf 100644 --- a/aiohttp_session/redis_storage.py +++ b/aiohttp_session/redis_storage.py @@ -11,12 +11,14 @@ class RedisStorage(AbstractStorage): def __init__(self, redis_pool, *, cookie_name="AIOHTTP_SESSION", domain=None, max_age=None, path='/', secure=None, httponly=True, - encoder=json.dumps, decoder=json.loads): + encoder=json.dumps, decoder=json.loads, + key_factory=lambda: uuid.uuid4().hex): super().__init__(cookie_name=cookie_name, domain=domain, max_age=max_age, path=path, secure=secure, httponly=httponly) self._encoder = encoder self._decoder = decoder + self._key_factory = key_factory self._redis = redis_pool @asyncio.coroutine @@ -42,7 +44,7 @@ def load_session(self, request): def save_session(self, request, response, session): key = session.identity if key is None: - key = uuid.uuid4().hex + key = self._key_factory() self.save_cookie(response, key, max_age=session.max_age) else: From f766413db3484314157dfd72b9aaa6a5536b5777 Mon Sep 17 00:00:00 2001 From: Smialy Date: Tue, 12 Sep 2017 21:42:12 +0200 Subject: [PATCH 2/2] Add key_factory unittest for redis_storage --- tests/test_redis_storage.py | 30 +++++++++++++++++++++++++++--- 1 file changed, 27 insertions(+), 3 deletions(-) diff --git a/tests/test_redis_storage.py b/tests/test_redis_storage.py index 8170863e..dfaab998 100644 --- a/tests/test_redis_storage.py +++ b/tests/test_redis_storage.py @@ -3,7 +3,6 @@ import uuid import aioredis import time - import pytest from aiohttp import web @@ -31,9 +30,10 @@ def start(): loop.run_until_complete(pool.clear()) -def create_app(loop, handler, redis, max_age=None): +def create_app(loop, handler, redis, max_age=None, + key_factory=lambda: uuid.uuid4().hex): middleware = session_middleware( - RedisStorage(redis, max_age=max_age)) + RedisStorage(redis, max_age=max_age, key_factory=key_factory)) app = web.Application(middlewares=[middleware], loop=loop) app.router.add_route('GET', '/', handler) return app @@ -262,3 +262,27 @@ def handler(request): {'AIOHTTP_SESSION': 'invalid_key'}) resp = yield from client.get('/') assert resp.status == 200 + + +@asyncio.coroutine +def test_create_storate_with_custom_key_factory(test_client, redis): + + @asyncio.coroutine + def handler(request): + session = yield from get_session(request) + session['key'] = 'value' + assert session.new + return web.Response(body=b'OK') + + def key_factory(): + return 'test-key' + + client = yield from test_client(create_app, handler, redis, 8, key_factory) + resp = yield from client.get('/') + assert resp.status == 200 + + assert resp.cookies['AIOHTTP_SESSION'].value == 'test-key' + + value = yield from load_cookie(client, redis) + assert 'key' in value['session'] + assert value['session']['key'] == 'value'