Fix #128

Use USERNAME_FIELD in utils.jwt_payload_handler

Author: jpadilla

rest_framework_jwt/authentication.py

@@ -2,13 +2,15 @@
 from django.utils.encoding import smart_text
 from django.utils.translation import ugettext as _
 from rest_framework import exceptions
-from rest_framework.authentication import (BaseAuthentication,
-                                           get_authorization_header)
+from rest_framework.authentication import (
+    BaseAuthentication, get_authorization_header
+)
 
-from rest_framework_jwt import utils
+from rest_framework_jwt.compat import get_user_model
 from rest_framework_jwt.settings import api_settings
 
 
+User = get_user_model()
 jwt_decode_handler = api_settings.JWT_DECODE_HANDLER
 jwt_get_user_id_from_payload = api_settings.JWT_PAYLOAD_GET_USER_ID_HANDLER
 
@@ -46,8 +48,6 @@ def authenticate_credentials(self, payload):
         """
         Returns an active user that matches the payload's user id and email.
         """
-        User = utils.get_user_model()
-
         user_id = jwt_get_user_id_from_payload(payload)
 
         if user_id is not None:

rest_framework_jwt/compat.py

@@ -1,4 +1,5 @@
 import rest_framework
+
 from distutils.version import StrictVersion
 
 
@@ -9,3 +10,32 @@ class Serializer(rest_framework.serializers.Serializer):
         @property
         def object(self):
             return self.validated_data
+
+
+def get_user_model():
+    try:
+        from django.contrib.auth import get_user_model
+    except ImportError:  # Django < 1.5
+        from django.contrib.auth.models import User
+    else:
+        User = get_user_model()
+
+    return User
+
+
+def get_username_field():
+    try:
+        username_field = get_user_model().USERNAME_FIELD
+    except:
+        username_field = 'username'
+
+    return username_field
+
+
+def get_username(user):
+    try:
+        username = user.get_username()
+    except AttributeError:
+        username = user.username
+
+    return username

rest_framework_jwt/serializers.py

@@ -8,10 +8,11 @@
 from rest_framework import serializers
 from .compat import Serializer
 
-from rest_framework_jwt import utils
+from rest_framework_jwt.compat import get_user_model, get_username_field
 from rest_framework_jwt.settings import api_settings
 
 
+User = get_user_model()
 jwt_payload_handler = api_settings.JWT_PAYLOAD_HANDLER
 jwt_encode_handler = api_settings.JWT_ENCODE_HANDLER
 jwt_decode_handler = api_settings.JWT_DECODE_HANDLER
@@ -39,12 +40,7 @@ def __init__(self, *args, **kwargs):
 
     @property
     def username_field(self):
-        User = utils.get_user_model()
-
-        try:
-            return User.USERNAME_FIELD
-        except AttributeError:
-            return 'username'
+        return get_username_field()
 
     def validate(self, attrs):
         credentials = {
@@ -107,7 +103,6 @@ def _check_payload(self, token):
         return payload
 
     def _check_user(self, payload):
-        User = utils.get_user_model()
         # Make sure user exists (may want to refactor this)
         try:
             user_id = jwt_get_user_id_from_payload(payload)

rest_framework_jwt/utils.py

@@ -2,33 +2,21 @@
 
 from datetime import datetime
 
+from rest_framework_jwt.compat import get_username, get_username_field
 from rest_framework_jwt.settings import api_settings
 
 
-def get_user_model():
-    try:
-        from django.contrib.auth import get_user_model
-    except ImportError:  # Django < 1.5
-        from django.contrib.auth.models import User
-    else:
-        User = get_user_model()
-
-    return User
-
-
 def jwt_payload_handler(user):
-    try:
-        username = user.get_username()
-    except AttributeError:
-        username = user.username
-
-    return {
+    payload = {
         'user_id': user.pk,
         'email': user.email,
-        'username': username,
         'exp': datetime.utcnow() + api_settings.JWT_EXPIRATION_DELTA
     }
 
+    payload[get_username_field()] = get_username(user)
+
+    return payload
+
 
 def jwt_get_user_id_from_payload_handler(payload):
     """

tests/test_authentication.py

@@ -2,7 +2,6 @@
 from django.test import TestCase
 from django.utils import unittest
 from django.conf.urls import patterns
-from django.contrib.auth import get_user_model
 
 from rest_framework import permissions, status
 try:
@@ -33,6 +32,7 @@
 from rest_framework.views import APIView
 
 from rest_framework_jwt import utils
+from rest_framework_jwt.compat import get_user_model
 from rest_framework_jwt.settings import api_settings, DEFAULTS
 from rest_framework_jwt.authentication import JSONWebTokenAuthentication
 

tests/test_serializers.py

@@ -3,8 +3,8 @@
 import rest_framework
 from django.test import TestCase
 from django.utils import unittest
-from django.contrib.auth import get_user_model
 
+from rest_framework_jwt.compat import get_user_model
 from rest_framework_jwt.serializers import JSONWebTokenSerializer
 from rest_framework_jwt import utils
 

tests/test_utils.py

@@ -1,10 +1,11 @@
 import json
 import base64
 
-from django.contrib.auth import get_user_model
-from django.test import TestCase
 import jwt.exceptions
+from django.test import TestCase
+
 from rest_framework_jwt import utils
+from rest_framework_jwt.compat import get_user_model
 from rest_framework_jwt.settings import api_settings, DEFAULTS
 
 User = get_user_model()

tests/test_views.py

@@ -6,14 +6,12 @@
 from django.test.utils import override_settings
 from django.utils import unittest
 from django.conf.urls import patterns
-from django.contrib.auth import get_user_model
-
 from freezegun import freeze_time
-
 from rest_framework import status
 from rest_framework.test import APIClient
 
 from rest_framework_jwt import utils
+from rest_framework_jwt.compat import get_user_model
 from rest_framework_jwt.settings import api_settings, DEFAULTS
 
 from . import utils as test_utils
@@ -192,7 +190,7 @@ def test_jwt_login_json(self):
 
         self.assertEqual(response.status_code, status.HTTP_200_OK)
         decoded_payload = utils.jwt_decode_handler(response.data['token'])
-        self.assertEqual(decoded_payload['username'], self.email)
+        self.assertEqual(decoded_payload['email'], self.email)
 
     def test_jwt_login_json_bad_creds(self):
         """

tests/utils.py

@@ -1,3 +1,6 @@
+from rest_framework_jwt.compat import get_username
+
+
 def jwt_response_payload_handler(token, user=None, request=None):
     """
     Returns the response data for both the login and refresh views.
@@ -14,6 +17,6 @@ def jwt_response_payload_handler(token, user=None, request=None):
 
     """
     return {
-        'user': user.get_username(),
+        'user': get_username(user),
         'token': token
     }