1# -*- coding: utf-8 -*-
2import webapp2
3from webapp2_extras import sessions
4
5import test_base
6
7app = webapp2.WSGIApplication(config={
8    'webapp2_extras.sessions': {
9        'secret_key': 'my-super-secret',
10    },
11})
12
13
14class 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
218if __name__ == '__main__':
219    test_base.main()
220