diff --git a/flask_jwt/__init__.py b/flask_jwt/__init__.py index f864b78..dd5688b 100644 --- a/flask_jwt/__init__.py +++ b/flask_jwt/__init__.py @@ -50,7 +50,7 @@ def _default_jwt_payload_handler(identity): iat = datetime.utcnow() exp = iat + current_app.config.get('JWT_EXPIRATION_DELTA') nbf = iat + current_app.config.get('JWT_NOT_BEFORE_DELTA') - identity = getattr(identity, 'id') or identity['id'] + identity = getattr(identity, 'id', None) or identity['id'] return {'exp': exp, 'iat': iat, 'nbf': nbf, 'identity': identity} diff --git a/tests/conftest.py b/tests/conftest.py index fb87ce4..413834a 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -67,3 +67,35 @@ def protected(): @pytest.fixture(scope='function') def client(app): return app.test_client() + + +@pytest.fixture(scope='function') +def dictuserapp(jwt, dictuser): + app = Flask(__name__) + app.debug = True + app.config['SECRET_KEY'] = 'super-secret' + + @jwt.authentication_handler + def authenticate(username, password): + if username == dictuser['username'] and password == dictuser['password']: + return dictuser + return None + + @jwt.identity_handler + def load_user(payload): + if payload['identity'] == dictuser['id']: + return dictuser + + jwt.init_app(app) + + @app.route('/protected') + @flask_jwt.jwt_required() + def protected(): + return 'success' + + return app + + +@pytest.fixture(scope='function') +def dictuser(): + return {'id': 1, 'username': 'joe', 'password': 'pass'} diff --git a/tests/test_jwt.py b/tests/test_jwt.py index 2157003..4970f05 100644 --- a/tests/test_jwt.py +++ b/tests/test_jwt.py @@ -291,3 +291,21 @@ def custom_auth_request_handler(): with app.test_client() as c: resp, jdata = post_json(c, '/auth', {}) assert jdata == {'hello': 'world'} + + +def test_default_encode_handler_user_object(app, client, jwt, user): + with app.app_context(): + token = jwt.jwt_encode_callback(user) + + with client as c: + c.get('/protected', headers={'authorization': 'JWT ' + token.decode('utf-8')}) + assert flask_jwt.current_identity == user + + +def test_default_encode_handler_dictuser(dictuserapp, jwt, dictuser): + with dictuserapp.app_context(): + token = jwt.jwt_encode_callback(dictuser) + + with dictuserapp.test_client() as c: + c.get('/protected', headers={'authorization': 'JWT ' + token.decode('utf-8')}) + assert flask_jwt.current_identity == dictuser