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

Add Django manage command to create a DRF user Token
This commit is contained in:
Tom Christie 2017-07-10 10:35:05 +01:00 committed by GitHub
commit fbb3490989
4 changed files with 80 additions and 0 deletions

View File

@ -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))

View File

@ -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