diff --git a/flask_jwt/__init__.py b/flask_jwt/__init__.py index f864b78..684efd2 100644 --- a/flask_jwt/__init__.py +++ b/flask_jwt/__init__.py @@ -50,7 +50,10 @@ 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'] + if hasattr(identity, 'id'): + identity = getattr(identity, 'id') + else: + identity = identity['id'] return {'exp': exp, 'iat': iat, 'nbf': nbf, 'identity': identity} diff --git a/tests/test_jwt.py b/tests/test_jwt.py index 2157003..55b1aa3 100644 --- a/tests/test_jwt.py +++ b/tests/test_jwt.py @@ -291,3 +291,27 @@ def custom_auth_request_handler(): with app.test_client() as c: resp, jdata = post_json(c, '/auth', {}) assert jdata == {'hello': 'world'} + + +def test_object_vs_dict_handling_in_default_jwt_payload_handler(app): + """Test dicts do not cause Attribute error when looking for id key or field. + """ + from flask_jwt import _default_jwt_payload_handler + + class IDT(object): + def __init__(self): + self.id = "456" + self.name = "tony" + + identity1 = IDT() + + identity2 = {"id": "123456", "name": "bob"} + + result = _default_jwt_payload_handler(identity1) + assert "identity" in result + assert result["identity"] == identity1.id + + # This did explode with attribute error when a dict was used: + result = _default_jwt_payload_handler(identity2) + assert "identity" in result + assert result["identity"] == identity2['id']