1 # -*- coding: utf-8 -*- 2 import os 3 4 import webapp2 5 from webapp2_extras import users 6 7 import test_base 8 9 10 def 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 16 class 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 26 class 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 36 app = webapp2.WSGIApplication([ 37 ('/login_required', LoginRequiredHandler), 38 ('/admin_required', AdminRequiredHandler), 39 ]) 40 41 42 class TestUsers(test_base.BaseTestCase): 43 def test_login_required_allowed(self): 44 set_current_user('foo (at] bar.com', 'foo (at] 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 (at] bar.com', 'foo (at] 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 (at] bar.com', 'foo (at] 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 97 if __name__ == '__main__': 98 test_base.main() 99
