Support for UUID for User PK

PierreF · PierreF · commit 630e0c97aad6 · 2016-04-03T15:19:03.000+02:00
diff --git a/rest_framework_jwt/utils.py b/rest_framework_jwt/utils.py
@@ -1,4 +1,5 @@
 import jwt
+import uuid
 import warnings
 from calendar import timegm
 from datetime import datetime
@@ -23,6 +24,8 @@ def jwt_payload_handler(user):
         'username': username,
         'exp': datetime.utcnow() + api_settings.JWT_EXPIRATION_DELTA
     }
+    if isinstance(user.pk, uuid.UUID):
+        payload['user_id'] = str(user.pk)
 
     payload[username_field] = username
 
diff --git a/tests/models.py b/tests/models.py
@@ -1,3 +1,4 @@
+import uuid
 from django.db import models
 from django.contrib.auth.models import AbstractBaseUser, BaseUserManager
 
@@ -11,3 +12,15 @@ class CustomUser(AbstractBaseUser):
 
     class Meta:
         app_label = 'tests'
+
+
+class CustomUserUUID(AbstractBaseUser):
+    id = models.UUIDField(primary_key=True, default=uuid.uuid4, editable=False)
+    email = models.EmailField(max_length=255, unique=True)
+
+    objects = BaseUserManager()
+
+    USERNAME_FIELD = 'email'
+
+    class Meta:
+        app_label = 'tests'
diff --git a/tests/test_views.py b/tests/test_views.py
@@ -205,6 +205,51 @@ def test_jwt_login_json_bad_creds(self):
         self.assertEqual(response.status_code, 400)
 
 
+@override_settings(AUTH_USER_MODEL='tests.CustomUserUUID')
+class CustomUserUUIDObtainJSONWebTokenTests(TestCase):
+    """JSON Web Token Authentication"""
+    urls = 'tests.test_views'
+
+    def setUp(self):
+        from .models import CustomUserUUID
+
+        self.email = 'jpueblo@example.com'
+        self.password = 'password'
+        user = CustomUserUUID.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['user_id'], str(self.user.id))
+
+    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)
+
+
 class TokenTestCase(BaseTestCase):
     """
     Handlers for getting tokens from the API, or creating arbitrary ones.
