https://groups.google.com/forum/#!topic/python-tornado/l8slxrWHq-c
import tornado.testing import tornado.web import tornado.httpserver import motor import functools import pymongo def authenticate(method): @tornado.gen.coroutine @functools.wraps(method) def wrapper(self, *args, **kwargs): key = self.request.headers['Authorization'] query = {'token': key} cursor = self.application.db.users.admins result = yield cursor.find_one(query, {'_id': 1}) if not result: self.set_status(401) return self.write('Unauthorized') else: return method(self, *args, **kwargs) return wrapper class GetCurrentUser(tornado.web.RequestHandler): @tornado.gen.coroutine def get(self): current_user = yield self.get_current_user() print('Current user:', current_user) @tornado.gen.coroutine def get_current_user(self): key = self.request.headers['Authorization'] query = {'token': key} cursor = self.application.db.users.admins result = yield cursor.find_one(query, {'_id': 1}) if result: return result['_id'] else: return None class GetDecoratedUser(tornado.web.RequestHandler): @authenticate def get(self): print('Success') class TestCase(tornado.testing.AsyncHTTPTestCase): def setUp(self): db = pymongo.MongoClient() db.users.admins.insert({'_id': 'bob', 'token': '12345'}) super(TestCase, self).setUp() def tearDown(self): db = pymongo.MongoClient() db.users.admins.remove({}) super(TestCase, self).tearDown() def get_app(self): app = tornado.web.Application([ (r'/current_user', GetCurrentUser), (r'/decorator', GetDecoratedUser) ]) app.db = motor.MotorClient() return app def test_get_current_user(self): self.fetch('/current_user', headers={'Authorization': '12345'}) self.fetch('/current_user', headers={'Authorization': '54321'}) def test_decorator(self): self.fetch('/decorator', headers={'Authorization': '12345'}) self.fetch('/decorator', headers={'Authorization': '54321'})
No comments:
Post a Comment