Support beaker (and possibly other) sessions

miguelgrinberg · Jun 4, 2017 · 2625426 · 2625426
1 parent 8d3d729
commit 2625426
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 13 deletions.
diff --git a/flask_socketio/__init__.py b/flask_socketio/__init__.py
@@ -15,7 +15,7 @@
 
 import socketio
 import flask
-from flask import json as flask_json
+from flask import _request_ctx_stack, json as flask_json
 from werkzeug.debug import DebuggedApplication
 from werkzeug.serving import run_with_reloader
 
@@ -583,10 +583,11 @@ def _handle_event(self, handler, message, namespace, sid, *args):
             return '', 400
         app = self.server.environ[sid]['flask.app']
         with app.request_context(self.server.environ[sid]):
-            if 'saved_session' in self.server.environ[sid]:
-                self._copy_session(
-                    self.server.environ[sid]['saved_session'],
-                    flask.session)
+            if 'saved_session' not in self.server.environ[sid]:
+                self.server.environ[sid]['saved_session'] = \
+                    dict(flask.session)
+            _request_ctx_stack.top.session = \
+                self.server.environ[sid]['saved_session']
             flask.request.sid = sid
             flask.request.namespace = namespace
             flask.request.event = {'message': message, 'args': args}
@@ -602,11 +603,6 @@ def _handle_event(self, handler, message, namespace, sid, *args):
                     raise
                 type, value, traceback = sys.exc_info()
                 return err_handler(value)
-            if flask.session.modified and sid in self.server.environ:
-                self.server.environ[sid]['saved_session'] = {}
-                self._copy_session(
-                    flask.session,
-                    self.server.environ[sid]['saved_session'])
             return ret
 
     def _copy_session(self, src, dest):

diff --git a/test_socketio.py b/test_socketio.py
@@ -366,10 +366,11 @@ def test_session(self):
         client = socketio.test_client(app)
         client.get_received()
         client.send('echo this message back')
-        self.assertNotIn('saved_session', socketio.server.environ[client.sid])
+        self.assertEqual(socketio.server.environ[client.sid]['saved_session'],
+                         {})
         client.send('test session')
-        session = socketio.server.environ[client.sid]['saved_session']
-        self.assertEqual(session['a'], 'b')
+        self.assertEqual(socketio.server.environ[client.sid]['saved_session'],
+                         {'a': 'b'})
 
     def test_room(self):
         client1 = socketio.test_client(app)