diff --git a/rest_framework/authtoken/management/__init__.py b/rest_framework/authtoken/management/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/rest_framework/authtoken/management/commands/__init__.py b/rest_framework/authtoken/management/commands/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/rest_framework/authtoken/management/commands/drf_create_token.py b/rest_framework/authtoken/management/commands/drf_create_token.py new file mode 100644 index 000000000..417bdd780 --- /dev/null +++ b/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)) diff --git a/tests/test_authtoken.py b/tests/test_authtoken.py index 54ac1848d..6374d7141 100644 --- a/tests/test_authtoken.py +++ b/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