Management command to reset program enrollment data (#21221)

mgmt cmd to reset program enrollments data

lms/djangoapps/program_enrollments/management/commands/reset_enrollment_data.py

@@ -0,0 +1,63 @@
+"""
+Management command to remove enrollments and any related models created as
+a side effect of enrolling students.
+
+Intented for use in integration sandbox environments
+"""
+from __future__ import absolute_import
+
+import logging
+
+from django.core.management.base import BaseCommand, CommandError
+from django.db import transaction
+from six.moves import input
+
+from lms.djangoapps.program_enrollments.models import ProgramEnrollment
+from student.models import CourseEnrollment
+
+log = logging.getLogger(__name__)
+
+
+class Command(BaseCommand):
+    """
+    Deletes all enrollments and related data
+
+    Example usage:
+        $ ./manage.py lms reset_enrollment_data ca73b4af-676a-4bb3-a9a5-f6b5a3dedd,1c5f61b9-0be5-4a90-9ea5-582d5e066c
+    """
+    confirmation_prompt = "Type 'confirm' to continue with deletion\n"
+
+    def add_arguments(self, parser):
+        parser.add_argument(
+            'programs',
+            help='Comma separated list of programs to delete enrollments for'
+        )
+        parser.add_argument(
+            '--force',
+            action='store_true',
+            help='Skip manual confirmation step before deleting objects',
+        )
+
+    @transaction.atomic
+    def handle(self, *args, **options):
+        programs = options['programs'].split(',')
+
+        q1_count, deleted_course_enrollment_models = CourseEnrollment.objects.filter(
+            programcourseenrollment__program_enrollment__program_uuid__in=programs
+        ).delete()
+        q2_count, deleted_program_enrollment_models = ProgramEnrollment.objects.filter(
+            program_uuid__in=programs
+        ).delete()
+
+        log.info(
+            u'The following records will be deleted:\n%s\n%s\n',
+            deleted_course_enrollment_models,
+            deleted_program_enrollment_models,
+        )
+
+        if not options['force']:
+            confirmation = input(self.confirmation_prompt)
+            if confirmation != 'confirm':
+                raise CommandError('User confirmation required.  No records have been modified')
+
+        log.info(u'Deleting %s records...', q1_count + q2_count)

lms/djangoapps/program_enrollments/management/commands/tests/test_reset_enrollment_data.py

@@ -0,0 +1,107 @@
+"""
+Tests for the reset_enrollment_data management command.
+"""
+from __future__ import absolute_import
+
+import sys
+from contextlib import contextmanager
+from StringIO import StringIO
+from uuid import uuid4
+
+from django.core.management import call_command
+from django.core.management.base import CommandError
+from django.test import TestCase
+
+from lms.djangoapps.program_enrollments.management.commands import reset_enrollment_data
+from lms.djangoapps.program_enrollments.models import ProgramCourseEnrollment, ProgramEnrollment
+from lms.djangoapps.program_enrollments.tests.factories import ProgramCourseEnrollmentFactory, ProgramEnrollmentFactory
+from student.models import CourseEnrollment
+from student.tests.factories import UserFactory
+
+
+class TestResetEnrollmentData(TestCase):
+    """ Test reset_enrollment_data command """
+
+    @classmethod
+    def setUpClass(cls):
+        super(TestResetEnrollmentData, cls).setUpClass()
+        cls.command = reset_enrollment_data.Command()
+        cls.program_uuid = uuid4()
+
+    def setUp(self):
+        super(TestResetEnrollmentData, self).setUp()
+        self.user = UserFactory()
+
+    @contextmanager
+    def _replace_stdin(self, text):
+        orig = sys.stdin
+        sys.stdin = StringIO(text)
+        yield
+        sys.stdin = orig
+
+    def _create_program_and_course_enrollment(self, program_uuid, user):
+        program_enrollment = ProgramEnrollmentFactory(user=user, program_uuid=program_uuid)
+        ProgramCourseEnrollmentFactory(program_enrollment=program_enrollment)
+        return program_enrollment
+
+    def _validate_enrollments_count(self, n):
+        self.assertEqual(len(CourseEnrollment.objects.all()), n)
+        self.assertEqual(len(ProgramCourseEnrollment.objects.all()), n)
+        self.assertEqual(len(ProgramEnrollment.objects.all()), n)
+
+    def test_reset(self):
+        """ Validate enrollments with a user and waiting enrollments without a user are removed """
+        self._create_program_and_course_enrollment(self.program_uuid, self.user)
+        self._create_program_and_course_enrollment(self.program_uuid, None)
+
+        call_command(self.command, self.program_uuid, force=True)
+
+        self._validate_enrollments_count(0)
+
+    def test_reset_confirmation(self):
+        """ By default this command will require user input to confirm """
+        self._create_program_and_course_enrollment(self.program_uuid, self.user)
+
+        with self._replace_stdin('confirm'):
+            call_command(self.command, self.program_uuid)
+
+        self._validate_enrollments_count(0)
+
+    def test_reset_confirmation_failure(self):
+        """ Failing to confirm reset will result in no modifications """
+        self._create_program_and_course_enrollment(self.program_uuid, self.user)
+
+        with self.assertRaises(CommandError):
+            with self._replace_stdin('no'):
+                call_command(self.command, self.program_uuid)
+
+        self._validate_enrollments_count(1)
+
+    def test_reset_scope(self):
+        """ reset should only affect provided programs """
+        self._create_program_and_course_enrollment(self.program_uuid, self.user)
+
+        alt_program_uuid = uuid4()
+        alt_user = UserFactory()
+        self._create_program_and_course_enrollment(alt_program_uuid, alt_user)
+
+        call_command(self.command, self.program_uuid, force=True)
+
+        # enrollment with different uuid still exists
+        program_enrollment = ProgramEnrollment.objects.get(program_uuid=alt_program_uuid)
+        program_course_enrollment = ProgramCourseEnrollment.objects.get(program_enrollment=program_enrollment)
+        course_enrollment = program_course_enrollment.course_enrollment
+        self.assertIsNotNone(course_enrollment)
+
+        # other enrollments have been deleted
+        self._validate_enrollments_count(1)
+
+    def test_reset_multiple_programs(self):
+        self._create_program_and_course_enrollment(self.program_uuid, self.user)
+        alt_program_uuid = uuid4()
+        alt_user = UserFactory()
+        self._create_program_and_course_enrollment(alt_program_uuid, alt_user)
+
+        call_command(self.command, '{},{}'.format(self.program_uuid, alt_program_uuid), force=True)
+
+        self._validate_enrollments_count(0)