Merge pull request #5188 from andreagrandi/auth-token-cmd

Add Django manage command to create a DRF user Token

rest_framework/authtoken/management/commands/drf_create_token.py

@@ -0,0 +1,45 @@
+from django.contrib.auth import get_user_model
+from django.core.management.base import BaseCommand, CommandError
+from rest_framework.authtoken.models import Token
+
+
+UserModel = get_user_model()
+
+
+class Command(BaseCommand):
+    help = 'Create DRF Token for a given user'
+
+    def create_user_token(self, username, reset_token):
+        user = UserModel._default_manager.get_by_natural_key(username)
+
+        if reset_token:
+            Token.objects.filter(user=user).delete()
+
+        token = Token.objects.get_or_create(user=user)
+        return token[0]
+
+    def add_arguments(self, parser):
+        parser.add_argument('username', type=str, nargs='+')
+
+        parser.add_argument(
+            '-r',
+            '--reset',
+            action='store_true',
+            dest='reset_token',
+            default=False,
+            help='Reset existing User token and create a new one',
+        )
+
+    def handle(self, *args, **options):
+        username = options['username']
+        reset_token = options['reset_token']
+
+        try:
+            token = self.create_user_token(username, reset_token)
+        except UserModel.DoesNotExist:
+            raise CommandError(
+                'Cannot create the Token: user {0} does not exist'.format(
+                    username)
+            )
+        self.stdout.write(
+            'Generated token {0} for user {1}'.format(token.key, username))

tests/test_authtoken.py

@@ -4,6 +4,8 @@ from django.contrib.auth.models import User
 from django.test import TestCase
 
 from rest_framework.authtoken.admin import TokenAdmin
+from rest_framework.authtoken.management.commands.drf_create_token import \
+    Command as AuthTokenCommand
 from rest_framework.authtoken.models import Token
 from rest_framework.authtoken.serializers import AuthTokenSerializer
 from rest_framework.exceptions import ValidationError
@@ -33,3 +35,36 @@ class AuthTokenTests(TestCase):
         self.user.set_password(data['password'])
         self.user.save()
         assert AuthTokenSerializer(data=data).is_valid()
+
+
+class AuthTokenCommandTests(TestCase):
+
+    def setUp(self):
+        self.site = site
+        self.user = User.objects.create_user(username='test_user')
+
+    def test_command_create_user_token(self):
+        token = AuthTokenCommand().create_user_token(self.user.username, False)
+        assert token is not None
+        token_saved = Token.objects.first()
+        assert token.key == token_saved.key
+
+    def test_command_create_user_token_invalid_user(self):
+        with pytest.raises(User.DoesNotExist):
+            AuthTokenCommand().create_user_token('not_existing_user', False)
+
+    def test_command_reset_user_token(self):
+        AuthTokenCommand().create_user_token(self.user.username, False)
+        first_token_key = Token.objects.first().key
+        AuthTokenCommand().create_user_token(self.user.username, True)
+        second_token_key = Token.objects.first().key
+
+        assert first_token_key != second_token_key
+
+    def test_command_do_not_reset_user_token(self):
+        AuthTokenCommand().create_user_token(self.user.username, False)
+        first_token_key = Token.objects.first().key
+        AuthTokenCommand().create_user_token(self.user.username, False)
+        second_token_key = Token.objects.first().key
+
+        assert first_token_key == second_token_key