1 # -*- coding: utf-8 -*- 2 import webapp2 3 from webapp2_extras import sessions 4 5 import test_base 6 7 app = webapp2.WSGIApplication(config={ 8 'webapp2_extras.sessions': { 9 'secret_key': 'my-super-secret', 10 }, 11 }) 12 13 14 class TestSecureCookieSession(test_base.BaseTestCase): 15 factory = sessions.SecureCookieSessionFactory 16 17 def test_config(self): 18 app = webapp2.WSGIApplication() 19 req = webapp2.Request.blank('/') 20 req.app = app 21 self.assertRaises(Exception, sessions.SessionStore, req) 22 23 # Just to set a special config. 24 app = webapp2.WSGIApplication() 25 req = webapp2.Request.blank('/') 26 req.app = app 27 store = sessions.SessionStore(req, config={ 28 'secret_key': 'my-super-secret', 29 'cookie_name': 'foo' 30 }) 31 session = store.get_session(factory=self.factory) 32 session['bar'] = 'bar' 33 rsp = webapp2.Response() 34 store.save_sessions(rsp) 35 self.assertTrue(rsp.headers['Set-Cookie'].startswith('foo=')) 36 37 def test_get_save_session(self): 38 # Round 1 ------------------------------------------------------------- 39 40 req = webapp2.Request.blank('/') 41 req.app = app 42 store = sessions.SessionStore(req) 43 44 session = store.get_session(factory=self.factory) 45 46 rsp = webapp2.Response() 47 # Nothing changed, we want to test anyway. 48 store.save_sessions(rsp) 49 50 session['a'] = 'b' 51 session['c'] = 'd' 52 session['e'] = 'f' 53 54 store.save_sessions(rsp) 55 56 # Round 2 ------------------------------------------------------------- 57 58 cookies = rsp.headers.get('Set-Cookie') 59 req = webapp2.Request.blank('/', headers=[('Cookie', cookies)]) 60 req.app = app 61 store = sessions.SessionStore(req) 62 63 session = store.get_session(factory=self.factory) 64 self.assertEqual(session['a'], 'b') 65 self.assertEqual(session['c'], 'd') 66 self.assertEqual(session['e'], 'f') 67 68 session['g'] = 'h' 69 70 rsp = webapp2.Response() 71 store.save_sessions(rsp) 72 73 # Round 3 ------------------------------------------------------------- 74 75 cookies = rsp.headers.get('Set-Cookie') 76 req = webapp2.Request.blank('/', headers=[('Cookie', cookies)]) 77 req.app = app 78 store = sessions.SessionStore(req) 79 80 session = store.get_session(factory=self.factory) 81 self.assertEqual(session['a'], 'b') 82 self.assertEqual(session['c'], 'd') 83 self.assertEqual(session['e'], 'f') 84 self.assertEqual(session['g'], 'h') 85 86 self.assertRaises(KeyError, session.pop, 'foo') 87 88 def test_flashes(self): 89 90 # Round 1 ------------------------------------------------------------- 91 92 req = webapp2.Request.blank('/') 93 req.app = app 94 store = sessions.SessionStore(req) 95 96 session = store.get_session(factory=self.factory) 97 flashes = session.get_flashes() 98 self.assertEqual(flashes, []) 99 session.add_flash('foo') 100 101 rsp = webapp2.Response() 102 store.save_sessions(rsp) 103 104 # Round 2 ------------------------------------------------------------- 105 106 cookies = rsp.headers.get('Set-Cookie') 107 req = webapp2.Request.blank('/', headers=[('Cookie', cookies)]) 108 req.app = app 109 store = sessions.SessionStore(req) 110 111 session = store.get_session(factory=self.factory) 112 113 flashes = session.get_flashes() 114 self.assertEqual(flashes, [[u'foo', None]]) 115 116 flashes = session.get_flashes() 117 self.assertEqual(flashes, []) 118 119 session.add_flash('bar') 120 session.add_flash('baz', 'important') 121 122 rsp = webapp2.Response() 123 store.save_sessions(rsp) 124 125 # Round 3 ------------------------------------------------------------- 126 127 cookies = rsp.headers.get('Set-Cookie') 128 req = webapp2.Request.blank('/', headers=[('Cookie', cookies)]) 129 req.app = app 130 store = sessions.SessionStore(req) 131 132 session = store.get_session(factory=self.factory) 133 134 flashes = session.get_flashes() 135 self.assertEqual(flashes, [[u'bar', None], [u'baz', 'important']]) 136 137 flashes = session.get_flashes() 138 self.assertEqual(flashes, []) 139 140 rsp = webapp2.Response() 141 store.save_sessions(rsp) 142 143 # Round 4 ------------------------------------------------------------- 144 145 cookies = rsp.headers.get('Set-Cookie') 146 req = webapp2.Request.blank('/', headers=[('Cookie', cookies)]) 147 req.app = app 148 store = sessions.SessionStore(req) 149 150 session = store.get_session(factory=self.factory) 151 flashes = session.get_flashes() 152 self.assertEqual(flashes, []) 153 154 def test_set_secure_cookie(self): 155 156 rsp = webapp2.Response() 157 158 # Round 1 ------------------------------------------------------------- 159 160 req = webapp2.Request.blank('/') 161 req.app = app 162 store = sessions.SessionStore(req) 163 164 store.set_secure_cookie('foo', {'bar': 'baz'}) 165 store.save_sessions(rsp) 166 167 # Round 2 ------------------------------------------------------------- 168 169 cookies = rsp.headers.get('Set-Cookie') 170 req = webapp2.Request.blank('/', headers=[('Cookie', cookies)]) 171 req.app = app 172 store = sessions.SessionStore(req) 173 res = store.get_secure_cookie('foo') 174 self.assertEqual(res, {'bar': 'baz'}) 175 176 def test_set_session_store(self): 177 app = webapp2.WSGIApplication(config={ 178 'webapp2_extras.sessions': { 179 'secret_key': 'my-super-secret', 180 } 181 }) 182 req = webapp2.Request.blank('/') 183 req.app = app 184 store = sessions.SessionStore(req) 185 186 self.assertEqual(len(req.registry), 0) 187 sessions.set_store(store, request=req) 188 self.assertEqual(len(req.registry), 1) 189 s = sessions.get_store(request=req) 190 self.assertTrue(isinstance(s, sessions.SessionStore)) 191 192 def test_get_session_store(self): 193 app = webapp2.WSGIApplication(config={ 194 'webapp2_extras.sessions': { 195 'secret_key': 'my-super-secret', 196 } 197 }) 198 req = webapp2.Request.blank('/') 199 req.app = app 200 self.assertEqual(len(req.registry), 0) 201 s = sessions.get_store(request=req) 202 self.assertEqual(len(req.registry), 1) 203 self.assertTrue(isinstance(s, sessions.SessionStore)) 204 205 def test_not_implemented(self): 206 req = webapp2.Request.blank('/') 207 req.app = app 208 store = sessions.SessionStore(req) 209 210 f = sessions.BaseSessionFactory('foo', store) 211 self.assertRaises(NotImplementedError, f.get_session) 212 self.assertRaises(NotImplementedError, f.save_session, None) 213 214 f = sessions.CustomBackendSessionFactory('foo', store) 215 self.assertRaises(NotImplementedError, f._get_by_sid, None) 216 217 218 if __name__ == '__main__': 219 test_base.main() 220