1# -*- coding: utf-8 -*-
2import os
3
4import webapp2
5from webapp2_extras import users
6
7import test_base
8
9
10def set_current_user(email, user_id, is_admin=False):
11    os.environ['USER_EMAIL'] = email or ''
12    os.environ['USER_ID'] = user_id or ''
13    os.environ['USER_IS_ADMIN'] = '1' if is_admin else '0'
14
15
16class LoginRequiredHandler(webapp2.RequestHandler):
17    @users.login_required
18    def get(self):
19        self.response.write('You are logged in.')
20
21    @users.login_required
22    def post(self):
23        self.response.write('You are logged in.')
24
25
26class AdminRequiredHandler(webapp2.RequestHandler):
27    @users.admin_required
28    def get(self):
29        self.response.write('You are admin.')
30
31    @users.admin_required
32    def post(self):
33        self.response.write('You are admin.')
34
35
36app = webapp2.WSGIApplication([
37    ('/login_required', LoginRequiredHandler),
38    ('/admin_required', AdminRequiredHandler),
39])
40
41
42class TestUsers(test_base.BaseTestCase):
43    def test_login_required_allowed(self):
44        set_current_user('foo@bar.com', 'foo@bar.com')
45        req = webapp2.Request.blank('/login_required')
46
47        rsp = req.get_response(app)
48        self.assertEqual(rsp.status_int, 200)
49        self.assertEqual(rsp.body, 'You are logged in.')
50
51    def test_login_required_302(self):
52        req = webapp2.Request.blank('/login_required')
53
54        rsp = req.get_response(app)
55        self.assertEqual(rsp.status_int, 302)
56        self.assertEqual(rsp.headers.get('Location'),
57            'https://www.google.com/accounts/Login?continue=http%3A//localhost/login_required')
58
59    def test_login_required_post(self):
60        req = webapp2.Request.blank('/login_required')
61        req.method = 'POST'
62
63        rsp = req.get_response(app)
64        self.assertEqual(rsp.status_int, 400)
65
66    def test_admin_required_allowed(self):
67        set_current_user('foo@bar.com', 'foo@bar.com', is_admin=True)
68        req = webapp2.Request.blank('/admin_required')
69
70        rsp = req.get_response(app)
71        self.assertEqual(rsp.status_int, 200)
72        self.assertEqual(rsp.body, 'You are admin.')
73
74    def test_admin_required_not_admin(self):
75        set_current_user('foo@bar.com', 'foo@bar.com')
76        req = webapp2.Request.blank('/admin_required')
77
78        rsp = req.get_response(app)
79        self.assertEqual(rsp.status_int, 403)
80
81    def test_admin_required_302(self):
82        req = webapp2.Request.blank('/admin_required')
83
84        rsp = req.get_response(app)
85        self.assertEqual(rsp.status_int, 302)
86        self.assertEqual(rsp.headers.get('Location'),
87            'https://www.google.com/accounts/Login?continue=http%3A//localhost/admin_required')
88
89    def test_admin_required_post(self):
90        req = webapp2.Request.blank('/admin_required')
91        req.method = 'POST'
92
93        rsp = req.get_response(app)
94        self.assertEqual(rsp.status_int, 400)
95
96
97if __name__ == '__main__':
98    test_base.main()
99