From 471d7311bf800fc6428903ce80fc49023f2c523d Mon Sep 17 00:00:00 2001
From: uzairr <uzairr@yahoo.com>
Date: Fri, 20 Mar 2020 17:49:01 +0500
Subject: [PATCH] Update bulk un-enroll management command

PROD-1370
---
 .../management/commands/bulk_unenroll.py      | 54 +++++++++++++++----
 1 file changed, 45 insertions(+), 9 deletions(-)

diff --git a/common/djangoapps/student/management/commands/bulk_unenroll.py b/common/djangoapps/student/management/commands/bulk_unenroll.py
index b1b4f87d83..12d043a1e7 100644
--- a/common/djangoapps/student/management/commands/bulk_unenroll.py
+++ b/common/djangoapps/student/management/commands/bulk_unenroll.py
@@ -1,5 +1,6 @@
-
-
+"""
+Un-enroll Bulk users course wide as well as provided in csv
+"""
 import logging
 
 import unicodecsv
@@ -7,13 +8,15 @@ from django.core.exceptions import ObjectDoesNotExist
 from django.core.management.base import BaseCommand
 from opaque_keys import InvalidKeyError
 from opaque_keys.edx.keys import CourseKey
-from student.models import CourseEnrollment, User, BulkUnenrollConfiguration
+from student.models import CourseEnrollment, BulkUnenrollConfiguration
 
 logger = logging.getLogger(__name__)  # pylint: disable=invalid-name
 
 
 class Command(BaseCommand):
-
+    """
+    Management command to un-enroll course enrollments at once.
+    """
     help = """"
     Un-enroll bulk users from the courses.
     It expect that the data will be provided in a csv file format with
@@ -27,18 +30,32 @@ class Command(BaseCommand):
                             dest='csv_path',
                             required=False,
                             help='Path to CSV file.')
+        parser.add_argument('-c', '--course-id',
+                            metavar='course_id',
+                            dest='course-id',
+                            required=False,
+                            help='unenroll all users from provided course-id')
 
     def handle(self, *args, **options):
 
         csv_path = options['csv_path']
+        course_id = options['course-id']
+
+        if course_id:
+            self.unenroll_all_users(course_id=course_id)
+            return
+
         if csv_path:
             with open(csv_path, 'rb') as csv_file:
-                self.unenroll_users(csv_file)
+                self.unenroll_users_from_csv(csv_file)
         else:
             csv_file = BulkUnenrollConfiguration.current().csv_file
-            self.unenroll_users(csv_file)
+            self.unenroll_users_from_csv(csv_file)
 
-    def unenroll_users(self, csv_file):
+    def unenroll_users_from_csv(self, csv_file):
+        """
+        Un-enroll the given set of users provided in csv
+        """
         reader = list(unicodecsv.DictReader(csv_file))
         users_unenrolled = {}
         for row in reader:
@@ -64,9 +81,28 @@ class Command(BaseCommand):
                 msg = 'Enrollment for the user {} in course {} does not exist!'.format(username, course_key)
                 logger.info(msg)
 
-            except Exception as err:
-                msg = 'Error un-enrolling User {} from course {}: '.format(username, course_key, err)
+            except Exception as err:  # pylint: disable=W0703
+                msg = 'Error un-enrolling User {} from course {}: with error: {}'.format(username, course_key, err)
                 logger.error(msg, exc_info=True)
 
         logger.info("Following users have been unenrolled successfully from the following courses: {users_unenrolled}"
                     .format(users_unenrolled=["{}:{}".format(k, v) for k, v in users_unenrolled.items()]))
+
+    def unenroll_all_users(self, course_id):
+        """
+        Un-enroll all users from a given course
+        """
+        try:
+            course_id = CourseKey.from_string(course_id)
+        except InvalidKeyError:
+            msg = 'Invalid course id {}, skipping un-enrollement.'.format(course_id)
+            logger.warning(msg)
+            return
+
+        try:
+            updated_count = CourseEnrollment.objects.filter(course_id=course_id, is_active=True).update(is_active=False)
+            logger.info("{} users have been unenrolled successfully from the provided course: {}"
+                        .format(updated_count, course_id))
+        except Exception as err:  # pylint: disable=W0703
+            msg = 'Error un-enrolling Users from course {}: with error: {}'.format(course_id, err)
+            logger.error(msg, exc_info=True)