Merge pull request #11281 from edx/bderusha/rc-change-enrollment-cmd

Update change_enrollment management command to be more robust and tested

common/djangoapps/student/management/commands/change_enrollment.py

+""" Command line script to change user enrollments. """
+
+import logging
+
 from django.core.management.base import BaseCommand, CommandError
+from django.db import transaction
 from opaque_keys import InvalidKeyError
+from opaque_keys.edx.keys import CourseKey
+from opaque_keys.edx.locations import SlashSeparatedCourseKey
 from optparse import make_option
+
 from student.models import CourseEnrollment, User
 
-from opaque_keys.edx.keys import CourseKey
-from opaque_keys.edx.locations import SlashSeparatedCourseKey
+logger = logging.getLogger(__name__)  # pylint: disable=invalid-name
+
+
+class RollbackException(Exception):
+    """
+    Exception raised explicitly to cause a database transaction rollback.
+    """
+    pass
 
 
 class Command(BaseCommand):
@@ -15,17 +29,17 @@ class Command(BaseCommand):
 
     Example:
 
-        Change enrollment for user joe from audit to honor:
+        Change enrollment for users joe, frank, and bill from audit to honor:
 
           $ ... change_enrollment -u joe,frank,bill -c some/course/id --from audit --to honor
 
         Or
 
-          $ ... change_enrollment -u "joe@example.com,frank@example.com,bill@example.com" -c some/course/id --from audit --to honor
+          $ ... change_enrollment -e "joe@example.com,frank@example.com,bill@example.com" -c some/course/id --from audit --to honor
 
-        Change enrollment for all users in some/course/id from audit to honor
+        See what would have been changed from audit to honor without making that change
 
-          $ ... change_enrollment -c some/course/id --from audit --to honor
+          $ ... change_enrollment -u joe,frank,bill -c some/course/id --from audit --to honor -n
 
     """
 
@@ -40,11 +54,16 @@ class Command(BaseCommand):
                     dest='to_mode',
                     default=False,
                     help='move to this enrollment mode'),
-        make_option('-u', '--user',
-                    metavar='USER',
-                    dest='user',
+        make_option('-u', '--usernames',
+                    metavar='USERNAME',
+                    dest='username',
+                    default=False,
+                    help="Comma-separated list of usernames to move in the course"),
+        make_option('-e', '--emails',
+                    metavar='EMAIL',
+                    dest='email',
                     default=False,
-                    help="Comma-separated list of users to move in the course"),
+                    help="Comma-separated list of email addresses to move in the course"),
         make_option('-c', '--course',
                     metavar='COURSE_ID',
                     dest='course_id',
@@ -59,8 +78,11 @@ class Command(BaseCommand):
     )
 
     def handle(self, *args, **options):
+        error_users = []
+        success_users = []
+
         if not options['course_id']:
-            raise CommandError("You must specify a course id for this command")
+            raise CommandError('You must specify a course id for this command')
         if not options['from_mode'] or not options['to_mode']:
             raise CommandError('You must specify a "to" and "from" mode as parameters')
 
@@ -69,26 +91,55 @@ class Command(BaseCommand):
         except InvalidKeyError:
             course_key = SlashSeparatedCourseKey.from_deprecated_string(options['course_id'])
 
-        filter_args = dict(
+        enrollment_args = dict(
             course_id=course_key,
             mode=options['from_mode']
         )
-        if options['user']:
-            user_str = options['user']
-            for username in user_str.split(","):
-                if '@' in username:
-                    user = User.objects.get(email=username)
-                else:
-                    user = User.objects.get(username=username)
-                filter_args['user'] = user
-                enrollments = CourseEnrollment.objects.filter(**filter_args)
-                if options['noop']:
-                    print "Would have changed {num_enrollments} students from {from_mode} to {to_mode}".format(
-                        num_enrollments=enrollments.count(),
-                        from_mode=options['from_mode'],
-                        to_mode=options['to_mode']
-                    )
-                else:
+
+        if options['username']:
+            self.update_enrollments('username', enrollment_args, options, error_users, success_users)
+
+        if options['email']:
+            self.update_enrollments('email', enrollment_args, options, error_users, success_users)
+
+        self.report(error_users, success_users)
+
+    def update_enrollments(self, identifier, enrollment_args, options, error_users, success_users):
+        """ Update enrollments for a specific user identifier (email or username). """
+        users = options[identifier].split(",")
+        for identified_user in users:
+            logger.info(identified_user)
+            try:
+                user_args = {
+                    identifier: identified_user
+                }
+
+                enrollment_args['user'] = User.objects.get(**user_args)
+                enrollments = CourseEnrollment.objects.filter(**enrollment_args)
+
+                with transaction.atomic():
                     for enrollment in enrollments:
                         enrollment.update_enrollment(mode=options['to_mode'])
                         enrollment.save()
+
+                    if options['noop']:
+                        raise RollbackException('Forced rollback.')
+
+            except RollbackException:
+                success_users.append(identified_user)
+                continue
+            except Exception as exception:  # pylint: disable=broad-except
+                error_users.append((identified_user, exception))
+                continue
+
+            success_users.append(identified_user)
+            logger.info('Updated user [%s] to mode [%s]', identified_user, options['to_mode'])
+
+    def report(self, error_users, success_users):
+        """ Log and overview of the results of the command. """
+        total_users = len(success_users) + len(error_users)
+        logger.info('Successfully updated %i out of %i users', len(success_users), total_users)
+        if len(error_users) > 0:
+            logger.info('The following %i user(s) not saved:', len(error_users))
+            for user, error in error_users:
+                logger.info('user: [%s] reason: [%s] %s', user, type(error).__name__, error.message)

common/djangoapps/student/management/tests/test_change_enrollment.py

+""" Test the change_enrollment command line script."""
+
+import ddt
+from mock import patch
+
+from django.core.management import call_command
+from xmodule.modulestore.tests.factories import CourseFactory
+from xmodule.modulestore.tests.django_utils import SharedModuleStoreTestCase
+
+from student.tests.factories import UserFactory, CourseModeFactory
+from student.models import CourseEnrollment
+
+
+@ddt.ddt
+class ChangeEnrollmentTests(SharedModuleStoreTestCase):
+    """ Test the enrollment change functionality of the change_enrollment script."""
+    def setUp(self):
+        super(ChangeEnrollmentTests, self).setUp()
+        self.course = CourseFactory.create()
+        self.audit_mode = CourseModeFactory.create(
+            course_id=self.course.id,
+            mode_slug='audit',
+            mode_display_name='Audit',
+        )
+        self.honor_mode = CourseModeFactory.create(
+            course_id=self.course.id,
+            mode_slug='honor',
+            mode_display_name='Honor',
+        )
+
+        self.user_info = [
+            ('amy', 'amy@pond.com', 'password'),
+            ('rory', 'rory@theroman.com', 'password'),
+            ('river', 'river@song.com', 'password')
+        ]
+        self.enrollments = []
+        self.users = []
+        for username, email, password in self.user_info:
+            user = UserFactory.create(username=username, email=email, password=password)
+            self.users.append(user)
+            self.enrollments.append(CourseEnrollment.enroll(user, self.course.id, mode='audit'))
+
+    @patch('student.management.commands.change_enrollment.logger')
+    @ddt.data(
+        ('email', False, 3),
+        ('username', False, 3),
+        ('email', True, 0),
+        ('username', True, 0),
+    )
+    @ddt.unpack
+    def test_convert_users(self, method, noop, expected_conversions, mock_logger):
+        """ The command should update the user's enrollment. """
+        user_str = ','.join([getattr(user, method) for user in self.users])
+        user_ids = [u.id for u in self.users]
+        command_args = {
+            'course_id': unicode(self.course.id),
+            'to_mode': 'honor',
+            'from_mode': 'audit',
+            'noop': noop,
+            method: user_str,
+        }
+
+        # Verify users are not in honor mode yet
+        self.assertEqual(
+            len(CourseEnrollment.objects.filter(mode='honor', user_id__in=user_ids)),
+            0
+        )
+
+        call_command(
+            'change_enrollment',
+            **command_args
+        )
+
+        # Verify correct number of users are now in honor mode
+        self.assertEqual(
+            len(CourseEnrollment.objects.filter(mode='honor', user_id__in=user_ids)),
+            expected_conversions
+        )
+
+        mock_logger.info.assert_called_with(
+            'Successfully updated %i out of %i users',
+            len(self.users),
+            len(self.users)
+        )
+
+    @patch('student.management.commands.change_enrollment.logger')
+    @ddt.data(
+        ('email', 'dtennant@thedoctor.com', 3),
+        ('username', 'dtennant', 3),
+    )
+    @ddt.unpack
+    def test_user_not_found(self, method, fake_user, expected_success, mock_logger):
+        all_users = [getattr(user, method) for user in self.users]
+        all_users.append(fake_user)
+        user_str = ','.join(all_users)
+        real_user_ids = [u.id for u in self.users]
+        command_args = {
+            'course_id': unicode(self.course.id),
+            'to_mode': 'honor',
+            'from_mode': 'audit',
+            method: user_str,
+        }
+
+        # Verify users are not in honor mode yet
+        self.assertEqual(
+            len(CourseEnrollment.objects.filter(mode='honor', user_id__in=real_user_ids)),
+            0
+        )
+
+        call_command(
+            'change_enrollment',
+            **command_args
+        )
+
+        # Verify correct number of users are now in honor mode
+        self.assertEqual(
+            len(CourseEnrollment.objects.filter(mode='honor', user_id__in=real_user_ids)),
+            expected_success
+        )
+
+        mock_logger.info.assert_called_with(
+            'user: [%s] reason: [%s] %s', fake_user, 'DoesNotExist', 'User matching query does not exist.'
+        )