Fix team size validation issue (#24290)

* Edit team manage to check sizes at end of import

* Fix size validation to take new teams into account

* Remove redundant max size check

* Consolidate team membership counters

* Remove unused user_ids_by_teamset_id set

* Fix team removal to only occur after validation

* Update team full error message

* Prefetch users when looking up team counts

lms/djangoapps/teams/csv.py

@@ -150,14 +150,12 @@ class TeamMembershipImportManager(object):
     def __init__(self, course):
         self.validation_errors = []
         self.teamset_ids = []
-        self.user_ids_by_teamset_id = {}
         self.course = course
         self.max_errors = 0
         self.existing_course_team_memberships = {}
         self.existing_course_teams = {}
         self.user_count_by_team = Counter()
         self.user_enrollment_by_team = {}
-        self.user_to_remove_by_team = Counter()
         self.number_of_learners_assigned = 0
         self.user_to_actual_enrollment_mode = {}
 
@@ -170,22 +168,32 @@ class TeamMembershipImportManager(object):
 
     def set_team_membership_from_csv(self, input_file):
         """
-        Assigns team membership based on the content of an uploaded CSV file.
+        Parse an input CSV file and pass to `set_team_memberships` for processing
+        """
+        csv_reader = csv.DictReader((line.decode('utf-8-sig').strip() for line in input_file.readlines()))
+        return self.set_team_memberships(csv_reader)
+
+    def set_team_memberships(self, csv_reader):
+        """
+        Assigns team membership based on the data from an uploaded CSV file.
         Returns true if there were no issues.
         """
-        reader = csv.DictReader((line.decode('utf-8-sig').strip() for line in input_file.readlines()))
-        self.teamset_ids = reader.fieldnames[2:]
+        # File-level validation
+        if not self.validate_header(csv_reader):
+            return False
+        if not self.validate_teamsets(csv_reader):
+            return False
+
+        self.teamset_ids = csv_reader.fieldnames[2:]
         row_dictionaries = []
         csv_usernames = set()
-        if not self.validate_header(reader.fieldnames):
-            return False
-        if not self.validate_teamsets():
-            return False
-        self.load_user_ids_by_teamset_id()
+
+        # Get existing team membership data
         self.load_course_team_memberships()
         self.load_course_teams()
+
         # process student rows:
-        for row in reader:
+        for row in csv_reader:
             if not self.validate_teams_have_matching_teamsets(row):
                 return False
             username = row['user']
@@ -203,11 +211,14 @@ class TeamMembershipImportManager(object):
             row['user'] = user
             if not self.validate_user_assignment_to_team_and_teamset(row):
                 return False
-            self.remove_user_from_team_for_reassignment(row)
             row_dictionaries.append(row)
 
+        if not self.validate_team_sizes_not_exceeded():
+            return False
+
         if not self.validation_errors:
             for row in row_dictionaries:
+                self.remove_user_from_team_for_reassignment(row)
                 self.add_user_to_team(row)
             self.number_of_learners_assigned = len(row_dictionaries)
             return True
@@ -226,15 +237,18 @@ class TeamMembershipImportManager(object):
     def load_course_teams(self):
         """
         Caches existing course teams by (team_name, topic_id)
+        and existing membership counts by (topic_id, team_name)
         """
-        for team in CourseTeam.objects.filter(course_id=self.course.id):
+        for team in CourseTeam.objects.filter(course_id=self.course.id).prefetch_related('users'):
             self.existing_course_teams[(team.name, team.topic_id)] = team
+            self.user_count_by_team[(team.topic_id, team.name)] = team.users.count()
 
-    def validate_header(self, header):
+    def validate_header(self, csv_reader):
         """
         Validates header row to ensure that it contains at a minimum columns called 'user', 'mode'.
         Teamset validation is handled separately
         """
+        header = csv_reader.fieldnames
         if 'user' not in header:
             self.validation_errors.append("Header must contain column 'user'.")
             return False
@@ -243,16 +257,18 @@ class TeamMembershipImportManager(object):
             return False
         return True
 
-    def validate_teamsets(self):
+    def validate_teamsets(self, csv_reader):
         """
         Validates team set ids. Returns true if there are no errors.
         The following conditions result in errors:
         Teamset does not exist
         Teamset id is duplicated
         """
+        teamset_ids = csv_reader.fieldnames[2:]
         valid_teamset_ids = {ts.teamset_id for ts in self.course.teams_configuration.teamsets}
+
         dupe_set = set()
-        for teamset_id in self.teamset_ids:
+        for teamset_id in teamset_ids:
             if teamset_id in dupe_set:
                 self.validation_errors.append("Teamset with id " + teamset_id + " is duplicated.")
                 return False
@@ -262,19 +278,6 @@ class TeamMembershipImportManager(object):
                 return False
         return True
 
-    def load_user_ids_by_teamset_id(self):
-        """
-        Get users associations with each teamset in a course and
-        save to `self.user_ids_by_teamset_id`
-        """
-        for teamset_id in self.teamset_ids:
-            self.user_ids_by_teamset_id[teamset_id] = {
-                membership.user_id for membership in
-                CourseTeamMembership.objects.filter(
-                    team__course_id=self.course.id, team__topic_id=teamset_id
-                )
-            }
-
     def validate_user_enrollment_is_valid(self, user, supplied_enrollment):
         """
         Invalid states:
@@ -327,33 +330,36 @@ class TeamMembershipImportManager(object):
         """
         user = row['user']
         for teamset_id in self.teamset_ids:
-            team_name = row[teamset_id]
-            if not team_name:
-                # we will be removing a user from their current team in a given teamset
-                # decrement the team membership, if it exists
-                try:
-                    current_team_name = self.existing_course_team_memberships[(user.id, teamset_id)].name
-                    if (teamset_id, current_team_name) not in self.user_to_remove_by_team:
-                        self.user_to_remove_by_team[(teamset_id, current_team_name)] = 1
-                    else:
-                        self.user_to_remove_by_team[(teamset_id, current_team_name)] += 1
-                except KeyError:
-                    # happens when a user is not on any team in a team set
-                    pass
-                continue
-            try:
-                if not self.validate_compatible_enrollment_modes(user, team_name, teamset_id):
-                    return False
-                # checks for a team inside a specific team set. This way team names can be duplicated across
-                # teamsets
-                team = self.existing_course_teams[(team_name, teamset_id)]
-                if (teamset_id, team_name) not in self.user_count_by_team:
-                    self.user_count_by_team[(teamset_id, team_name)] = team.users.count()
-            except KeyError:
-                pass
+            # See if the user is already on a team in the teamset
+            if (user.id, teamset_id) in self.existing_course_team_memberships:
+                current_team_name = self.existing_course_team_memberships[(user.id, teamset_id)].name
+            else:
+                current_team_name = None
 
-            if not self.validate_proposed_team_size_wont_exceed_maximum(team_name, teamset_id):
+            team_name = row[teamset_id]
+
+            # We don't need to do anything if the user isn't moving to a different team
+            if current_team_name == team_name:
+                continue
+
+            # If the user is on a team currently, remove them in from the updated count
+            if current_team_name is not None:
+                self.user_count_by_team[(teamset_id, current_team_name)] -= 1
+
+            # If we aren't moving them to a new team, we can go to the next team-set
+            if not team_name:
+                continue
+
+            # Check that user enrollment mode is compatible for the target team
+            if not self.validate_compatible_enrollment_modes(user, team_name, teamset_id):
                 return False
+
+            # Update proposed team counts, initializing the team count if it doesn't exist
+            if (teamset_id, team_name) not in self.user_count_by_team:
+                self.user_count_by_team[(teamset_id, team_name)] = 1
+            else:
+                self.user_count_by_team[(teamset_id, team_name)] += 1
+
         return True
 
     def validate_compatible_enrollment_modes(self, user, team_name, teamset_id):
@@ -406,21 +412,34 @@ class TeamMembershipImportManager(object):
         else:
             return True
 
-    def validate_proposed_team_size_wont_exceed_maximum(self, team_name, teamset_id):
+    def validate_team_sizes_not_exceeded(self):
         """
         Validates that the number of users we want to add to a team won't exceed maximum team size.
         """
-        if self.course.teams_configuration.teamsets_by_id[teamset_id].max_team_size is None:
-            max_team_size = self.course.teams_configuration.default_max_team_size
-        else:
-            max_team_size = self.course.teams_configuration.teamsets_by_id[teamset_id].max_team_size
+        for teamset_id in self.teamset_ids:
+            # Get max size for team-set
+            if self.course.teams_configuration.teamsets_by_id[teamset_id].max_team_size is None:
+                max_team_size = self.course.teams_configuration.default_max_team_size
+            else:
+                max_team_size = self.course.teams_configuration.teamsets_by_id[teamset_id].max_team_size
+
+            # Get teams in team-set
+            team_names = [
+                teamset_to_team[1] for teamset_to_team in self.user_count_by_team
+                if teamset_to_team[0] == teamset_id
+            ]
+
+            # Calculate proposed team size and return False if it exceeds capacity
+            for team_name in team_names:
+                key = (teamset_id, team_name)
+                if self.user_count_by_team[key] > max_team_size:
+                    self.add_error_and_check_if_max_exceeded(
+                        'New membership for team {} would exceed max size of {}.'.format(
+                            team_name, max_team_size
+                        )
+                    )
+                    return False
 
-        if max_team_size is not None:
-            key = (teamset_id, team_name)
-            if (self.user_count_by_team[key] - self.user_to_remove_by_team[key]) + 1 > max_team_size:
-                self.add_error_and_check_if_max_exceeded('Team ' + team_name + ' is full.')
-                return False
-        self.user_count_by_team[(teamset_id, team_name)] += 1
         return True
 
     def remove_user_from_team_for_reassignment(self, row):
@@ -447,7 +466,6 @@ class TeamMembershipImportManager(object):
                         try:
                             self._remove_user_from_teamset_and_emit_signal(row['user'].id, ts_id, self.course.id)
                             del self.existing_course_team_memberships[row['user'].id, ts_id]
-                            self.user_ids_by_teamset_id[ts_id].remove(row['user'].id)
                         except CourseTeamMembership.DoesNotExist:
                             pass
                     else:

lms/djangoapps/teams/tests/test_csv.py

@@ -2,6 +2,8 @@
 from csv import DictWriter, DictReader
 from io import BytesIO, StringIO, TextIOWrapper
 
+from django.contrib.auth.models import User
+
 from lms.djangoapps.program_enrollments.tests.factories import ProgramEnrollmentFactory, ProgramCourseEnrollmentFactory
 from lms.djangoapps.teams import csv
 from lms.djangoapps.teams.models import CourseTeam, CourseTeamMembership
@@ -241,14 +243,31 @@ class TeamMembershipImportManagerTests(TeamMembershipEventTestMixin, SharedModul
                 'id': 'teamset_1',
                 'name': 'teamset_name',
                 'description': 'teamset_desc',
+                'max_team_size': 3,
             }]
         })
         cls.course = CourseFactory(teams_configuration=teams_config)
         cls.second_course = CourseFactory(teams_configuration=teams_config)
+        cls.import_manager = None
 
-        # initialize import manager
-        cls.import_manager = csv.TeamMembershipImportManager(cls.course)
-        cls.import_manager.teamset_ids = {ts.teamset_id for ts in cls.course.teamsets}
+    def setUp(self):
+        """ Initialize import manager """
+        super(TeamMembershipImportManagerTests, self).setUp()
+        self.import_manager = csv.TeamMembershipImportManager(self.course)
+        self.import_manager.teamset_ids = {ts.teamset_id for ts in self.course.teamsets}
+
+    def test_load_course_teams(self):
+        """
+        Lodaing course teams shold get the users by team with only 2 queries
+        1 for teams, 1 for user count
+        """
+        team1 = CourseTeamFactory.create(course_id=self.course.id)
+        team2 = CourseTeamFactory.create(course_id=self.course.id)
+        team3 = CourseTeamFactory.create(course_id=self.course.id)
+        team4 = CourseTeamFactory.create(course_id=self.course.id)
+
+        with self.assertNumQueries(2):
+            self.import_manager.load_course_teams()
 
     def test_add_user_to_new_protected_team(self):
         """Adding a masters learner to a new team should create a team with organization protected status"""
@@ -330,6 +349,134 @@ class TeamMembershipImportManagerTests(TeamMembershipEventTestMixin, SharedModul
         self.assert_learner_removed_emitted(team_1.team_id, audit_learner.id)
         self.assert_learner_added_emitted(team_2.team_id, audit_learner.id)
 
+    def test_exceed_max_size(self):
+        # Given a bunch of students enrolled in a course
+        users = []
+        for i in range(5):
+            user = UserFactory.create(username='max_size_{id}'.format(id=i))
+            CourseEnrollmentFactory.create(user=user, course_id=self.course.id, mode='audit')
+            users.append(user)
+
+        # When a team is already near capaciy
+        team = CourseTeam.objects.create(
+            name='team_1',
+            course_id=self.course.id,
+            topic_id='teamset_1',
+            description='Team 1!',
+        )
+        for i in range(2):
+            user = users[i]
+            team.add_user(user)
+
+        # ... and I try to add members in excess of capacity
+        csv_data = self._csv_reader_from_array([
+            ['user', 'mode', 'teamset_1'],
+            ['max_size_0', 'audit', ''],
+            ['max_size_2', 'audit', 'team_1'],
+            ['max_size_3', 'audit', 'team_1'],
+            ['max_size_4', 'audit', 'team_1'],
+        ])
+
+        result = self.import_manager.set_team_memberships(csv_data)
+
+        # Then the import fails with no events emitted and a "team is full" error
+        self.assertFalse(result)
+        self.assert_no_events_were_emitted()
+        self.assertEqual(
+            self.import_manager.validation_errors[0],
+            'New membership for team team_1 would exceed max size of 3.'
+        )
+
+        # Confirm that memberships were not altered
+        for i in range(2):
+            self.assertTrue(CourseTeamMembership.is_user_on_team(user, team))
+
+    def test_remove_from_team(self):
+        # Given a user already in a course and on a team
+        user = UserFactory.create(username='learner_1')
+        mode = 'audit'
+        CourseEnrollmentFactory.create(user=user, course_id=self.course.id, mode=mode)
+        team = CourseTeamFactory(course_id=self.course.id, name='team_1', topic_id='teamset_1')
+        team.add_user(user)
+        self.assertTrue(CourseTeamMembership.is_user_on_team(user, team))
+
+        # When I try to remove them from the team
+        csv_data = self._csv_reader_from_array([
+            ['user', 'mode', 'teamset_1'],
+            [user.username, mode, ''],
+        ])
+        result = self.import_manager.set_team_memberships(csv_data)
+
+        # Then they are removed from the team and the correct events are issued
+        self.assertFalse(CourseTeamMembership.is_user_on_team(user, team))
+        self.assert_learner_removed_emitted(team.team_id, user.id)
+
+    def test_switch_memberships(self):
+        # Given a bunch of students enrolled in a course
+        users = []
+        for i in range(5):
+            user = UserFactory.create(username='learner_{id}'.format(id=i))
+            CourseEnrollmentFactory.create(user=user, course_id=self.course.id, mode='audit')
+            users.append(user)
+
+        # When a team is already at/near capaciy
+        for i in range(3):
+            user = users[i]
+            row = {'user': user, 'teamset_1': 'team_1', 'mode': 'audit'}
+            self.import_manager.add_user_to_team(row)
+
+        # ... and I try to switch membership (add/remove)
+        csv_data = self._csv_reader_from_array([
+            ['user', 'mode', 'teamset_1'],
+            ['learner_4', 'audit', 'team_1'],
+            ['learner_0', 'audit', 'team_2'],
+        ])
+
+        result = self.import_manager.set_team_memberships(csv_data)
+
+        # Then membership size is calculated correctly, import finishes w/out error
+        self.assertTrue(result)
+
+        # ... and the users are assigned to the correct teams
+        team_1 = CourseTeam.objects.get(course_id=self.course.id, topic_id='teamset_1', name='team_1')
+        self.assertTrue(CourseTeamMembership.is_user_on_team(users[4], team_1))
+        self.assert_learner_added_emitted(team_1.team_id, users[4].id)
+
+        team_2 = CourseTeam.objects.get(course_id=self.course.id, topic_id='teamset_1', name='team_2')
+        self.assertTrue(CourseTeamMembership.is_user_on_team(users[0], team_2))
+        self.assert_learner_added_emitted(team_2.team_id, users[0].id)
+
+    def test_create_new_team_from_import(self):
+        # Given a user in a course
+        user = UserFactory.create(username='learner_1')
+        mode = 'audit'
+        CourseEnrollmentFactory.create(user=user, course_id=self.course.id, mode=mode)
+
+        # When I add them to a team that does not exist
+        self.assertEquals(CourseTeam.objects.all().count(), 0)
+        csv_data = self._csv_reader_from_array([
+            ['user', 'mode', 'teamset_1'],
+            [user.username, mode, 'new_exciting_team'],
+        ])
+        result = self.import_manager.set_team_memberships(csv_data)
+
+        # Then a new team is created
+        self.assertEqual(CourseTeam.objects.all().count(), 1)
+
+        # ... and the user is assigned to the team
+        new_team = CourseTeam.objects.get(topic_id='teamset_1', name='new_exciting_team')
+        self.assertTrue(CourseTeamMembership.is_user_on_team(user, new_team))
+        self.assert_learner_added_emitted(new_team.team_id, user.id)
+
+    def _csv_reader_from_array(self, rows):
+        """
+        Given a 2D array, treat each element as a cell of a CSV file and construct a reader
+
+        Example:
+            [['header1', 'header2'], ['r1:c1', 'r1:c2'], ['r2:c2', 'r3:c3'] ... ]
+        """
+        return DictReader((','.join(row) for row in rows))
+
 
 class ExternalKeyCsvTests(TeamMembershipEventTestMixin, SharedModuleStoreTestCase):
     """ Tests for functionality related to external_user_keys"""

lms/djangoapps/teams/tests/test_views.py

@@ -2926,7 +2926,10 @@ class TestBulkMembershipManagement(TeamAPITestCase):
             data={'csv': csv_file}, user='staff'
         )
         response_text = json.loads(response.content.decode('utf-8'))
-        self.assertEqual(response_text['errors'][0], 'Team {} is full.'.format(team1))
+        self.assertEqual(
+            response_text['errors'][0],
+            'New membership for team {} would exceed max size of {}.'.format(team1, 3)
+        )
 
     def test_deletion_via_upload_csv(self):
         # create a team membership that will be used further down