Fix with test: _default_jwt_payload_handler doesn't handle dicts as the code seems to imply it should

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}
 
 

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']