Merge pull request #21 from doordash/custom-user-model

Custom User Models + Tests

Author: jpadilla

rest_framework_jwt/runtests/models.py

@@ -0,0 +1,5 @@
+from django.contrib.auth.models import User
+
+
+class CustomUser(User):
+    USERNAME_FIELD = 'email'

rest_framework_jwt/runtests/settings.py

@@ -21,6 +21,7 @@
     'django.contrib.sessions',
     'django.contrib.messages',
     'django.contrib.staticfiles',
+    'rest_framework_jwt.runtests',
 )
 
 # OAuth2 is optional and won't work if there is no provider & oauth2

rest_framework_jwt/serializers.py

@@ -1,4 +1,4 @@
-from django.contrib.auth import authenticate
+from django.contrib.auth import authenticate, get_user_model
 from rest_framework import serializers
 
 from rest_framework_jwt.settings import api_settings
@@ -12,17 +12,27 @@ class JSONWebTokenSerializer(serializers.Serializer):
     """
     Serializer class used to validate a username and password.
 
+    'username' is identified by the custom UserModel.USERNAME_FIELD.
+
     Returns a JSON Web Token that can be used to authenticate later calls.
     """
-    username = serializers.CharField()
+
     password = serializers.CharField(write_only=True)
 
-    def validate(self, attrs):
-        username = attrs.get('username')
-        password = attrs.get('password')
+    def __init__(self, *args, **kwargs):
+        """Dynamically add the USERNAME_FIELD to self.fields."""
+        super(JSONWebTokenSerializer, self).__init__(*args, **kwargs)
+        self.fields[self.username_field] = serializers.CharField()
 
-        if username and password:
-            user = authenticate(username=username, password=password)
+    @property
+    def username_field(self):
+        return get_user_model().USERNAME_FIELD
+
+    def validate(self, attrs):
+        credentials = {self.username_field: attrs.get(self.username_field),
+                       'password': attrs.get('password')}
+        if all(credentials.values()):
+            user = authenticate(**credentials)
 
             if user:
                 if not user.is_active:

rest_framework_jwt/tests/test_views.py

@@ -1,11 +1,12 @@
 from django.test import TestCase
+from django.conf import settings
 from django.contrib.auth.models import User
 from rest_framework import status
 from rest_framework.compat import patterns
 from rest_framework.test import APIClient
 
 from rest_framework_jwt import utils
-
+from rest_framework_jwt.runtests.models import CustomUser
 
 urlpatterns = patterns(
     '',
@@ -76,3 +77,52 @@ def test_jwt_login_form(self):
 
         self.assertEqual(response.status_code, status.HTTP_200_OK)
         self.assertEqual(decoded_payload['username'], self.username)
+
+
+class CustomUserObtainJSONWebTokenTests(TestCase):
+    """JSON Web Token Authentication"""
+    urls = 'rest_framework_jwt.tests.test_views'
+
+    def setUp(self):
+        # set custom user model
+        self.ORIG_AUTH_USER_MODEL = settings.AUTH_USER_MODEL
+        settings.AUTH_USER_MODEL = 'runtests.CustomUser'
+
+        self.email = '[email protected]'
+        self.password = 'password'
+        user = CustomUser.objects.create(email=self.email)
+        user.set_password(self.password)
+        user.save()
+        self.user = user
+
+        self.data = {
+            'email': self.email,
+            'password': self.password
+        }
+
+    def test_jwt_login_json(self):
+        """
+        Ensure JWT login view using JSON POST works.
+        """
+        client = APIClient(enforce_csrf_checks=True)
+
+        response = client.post('/auth-token/', self.data, format='json')
+
+        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)
+
+    def test_jwt_login_json_bad_creds(self):
+        """
+        Ensure JWT login view using JSON POST fails
+        if bad credentials are used.
+        """
+        client = APIClient(enforce_csrf_checks=True)
+
+        self.data['password'] = 'wrong'
+        response = client.post('/auth-token/', self.data, format='json')
+
+        self.assertEqual(response.status_code, 400)
+
+    def tearDown(self):
+        settings.AUTH_USER_MODEL = self.ORIG_AUTH_USER_MODEL