Added --users option to waffle_flag management command

Administrators may pass the `--users` option to activate a flag for a set of users, identified by their usernames and/or email addresses.

Author: acraftydev

.gitignore

@@ -6,6 +6,7 @@ dist
 *test.db
 .coverage
 .idea/
+*.iml
 .project
 .pydevproject
 .settings

waffle/management/commands/waffle_flag.py

@@ -1,10 +1,14 @@
 from __future__ import print_function
 
+from django.contrib.auth import get_user_model
 from django.contrib.auth.models import Group
 from django.core.management.base import BaseCommand, CommandError
+from django.db.models import Q
 
 from waffle import get_waffle_flag_model
 
+UserModel = get_user_model()
+
 
 class Command(BaseCommand):
     def add_arguments(self, parser):
@@ -60,6 +64,13 @@ def add_arguments(self, parser):
             help='Turn on the flag for listed group names (use flag more '
                  'than once for multiple groups). WARNING: This will remove '
                  'any currently associated groups unless --append is used!')
+        parser.add_argument(
+            '--user', '-u',
+            action='append',
+            default=list(),
+            help='Turn on the flag for listed usernames (use flag more '
+                 'than once for multiple users). WARNING: This will remove '
+                 'any currently associated users unless --append is used!')
         parser.add_argument(
             '--append',
             action='store_true',
@@ -97,6 +108,9 @@ def handle(self, *args, **options):
                 self.stdout.write('GROUPS: %s' % list(
                     flag.groups.values_list('name', flat=True))
                 )
+                self.stdout.write('USERS: %s' % list(
+                    flag.users.values_list(UserModel.USERNAME_FIELD, flat=True))
+                                  )
                 self.stdout.write('')
             return
 
@@ -137,6 +151,23 @@ def handle(self, *args, **options):
                 )
                 for group_name, group_id in group_hash.items():
                     flag.groups.add(group_id)
+            elif option == 'user':
+                user_hash = set()
+                for username in options['user']:
+                    try:
+                        user_instance = UserModel.objects.get(
+                            Q(**{UserModel.USERNAME_FIELD: username}) |
+                            Q(**{UserModel.EMAIL_FIELD: username})
+                        )
+                        user_hash.add(user_instance)
+                    except UserModel.DoesNotExist:
+                        raise CommandError('User %s does not exist' % username)
+                # If 'append' was not passed, we clear related users
+                if not options['append']:
+                    flag.users.clear()
+                self.stdout.write('Setting user(s): %s' % user_hash)
+                # for user in user_hash:
+                flag.users.add(*[user.id for user in user_hash])
             elif hasattr(flag, option):
                 self.stdout.write('Setting %s: %s' % (option, options[option]))
                 setattr(flag, option, options[option])

waffle/tests/test_management.py

@@ -2,7 +2,7 @@
 
 import six
 from django.core.management import call_command, CommandError
-from django.contrib.auth.models import Group
+from django.contrib.auth.models import Group, User
 
 from waffle import get_waffle_flag_model
 from waffle.models import Sample, Switch
@@ -120,7 +120,7 @@ def test_list(self):
         call_command('waffle_flag', list_flags=True, stdout=stdout)
         expected = 'Flags:\nNAME: test\nSUPERUSERS: True\nEVERYONE: None\n' \
                    'AUTHENTICATED: False\nPERCENT: None\nTESTING: False\n' \
-                   'ROLLOUT: False\nSTAFF: False\nGROUPS: []'
+                   'ROLLOUT: False\nSTAFF: False\nGROUPS: []\nUSERS: []'
         actual = stdout.getvalue().strip()
         self.assertEqual(actual, expected)
 
@@ -143,6 +143,52 @@ def test_group_append(self):
                          ['waffle_group', 'append_group'])
         self.assertIsNone(flag.everyone)
 
+    def test_user(self):
+        """ The command should replace a user to a flag."""
+        original_user = User.objects.create_user('waffle_test')
+        User.objects.create_user('add_user')
+        flag = get_waffle_flag_model().objects.create(name='test')
+        flag.users.add(original_user)
+        flag.refresh_from_db()
+
+        self.assertEqual(list(flag.users.values_list('username', flat=True)),
+                         ['waffle_test'])
+
+        call_command('waffle_flag', 'test', user=['add_user'])
+
+        flag.refresh_from_db()
+        self.assertEqual(list(flag.users.values_list('username', flat=True)),
+                         ['add_user'])
+        self.assertIsNone(flag.everyone)
+
+    def test_user_append(self):
+        """ The command should append a user to a flag."""
+        original_user = User.objects.create_user('waffle_test')
+        User.objects.create_user('append_user')
+        User.objects.create_user('append_user_email', email='[email protected]')
+        flag = get_waffle_flag_model().objects.create(name='test')
+        flag.users.add(original_user)
+        flag.refresh_from_db()
+
+        self.assertEqual(list(flag.users.values_list('username', flat=True)),
+                         ['waffle_test'])
+
+        call_command('waffle_flag', 'test', user=['append_user'],
+                     append=True)
+
+        flag.refresh_from_db()
+        self.assertEqual(list(flag.users.values_list('username', flat=True)),
+                         ['waffle_test', 'append_user'])
+        self.assertIsNone(flag.everyone)
+
+        call_command('waffle_flag', 'test', user=['[email protected]'],
+                     append=True)
+
+        flag.refresh_from_db()
+        self.assertEqual(list(flag.users.values_list('username', flat=True)),
+                         ['waffle_test', 'append_user', 'append_user_email'])
+        self.assertIsNone(flag.everyone)
+
 
 class WaffleSampleManagementCommandTests(TestCase):
     def test_create(self):