From d2459710cae9e17b785e1d67b11d7bc9aa39678a Mon Sep 17 00:00:00 2001
From: Andrea Grandi <a.grandi@gmail.com>
Date: Sat, 3 Jun 2017 11:58:01 +0100
Subject: [PATCH] Implement option to reset User token

---
 .../management/commands/drf_create_token.py   | 18 +++++++++++++++--
 tests/test_authtoken.py                       | 20 +++++++++++++++++--
 2 files changed, 34 insertions(+), 4 deletions(-)

diff --git a/rest_framework/authtoken/management/commands/drf_create_token.py b/rest_framework/authtoken/management/commands/drf_create_token.py
index 8fe47c44a..417bdd780 100644
--- a/rest_framework/authtoken/management/commands/drf_create_token.py
+++ b/rest_framework/authtoken/management/commands/drf_create_token.py
@@ -9,19 +9,33 @@ UserModel = get_user_model()
 class Command(BaseCommand):
     help = 'Create DRF Token for a given user'
 
-    def create_user_token(self, username):
+    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)
+            token = self.create_user_token(username, reset_token)
         except UserModel.DoesNotExist:
             raise CommandError(
                 'Cannot create the Token: user {0} does not exist'.format(
diff --git a/tests/test_authtoken.py b/tests/test_authtoken.py
index 407d82944..6374d7141 100644
--- a/tests/test_authtoken.py
+++ b/tests/test_authtoken.py
@@ -44,11 +44,27 @@ class AuthTokenCommandTests(TestCase):
         self.user = User.objects.create_user(username='test_user')
 
     def test_command_create_user_token(self):
-        token = AuthTokenCommand().create_user_token(self.user.username)
+        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')
+            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