test: refactor tests

edx_exams/apps/core/management/commands/test/test_bulk_add_course_staff.py

@@ -36,6 +36,15 @@ def _write_test_csv(self, csv, lines):
         csv.seek(0)
         return csv
 
+    def _assert_user_and_role(self, username, email, course_id, course_role):
+        user = User.objects.filter(username=username, email=email)
+        assert user.exists()
+        assert CourseStaffRole.objects.filter(
+            user=user[0].id,
+            course_id=course_id,
+            role=course_role,
+        ).exists()
+
     def test_empty_csv(self):
         lines = []
         with NamedTemporaryFile() as csv:
@@ -47,12 +56,12 @@ def test_add_course_staff_with_existing_user(self):
         with NamedTemporaryFile() as csv:
             csv = self._write_test_csv(csv, lines)
             call_command(self.command, f'--csv_path={csv.name}')
-            assert User.objects.filter(username=self.user.username, email=self.user.email).exists()
-            assert CourseStaffRole.objects.filter(
-                user=self.user.id,
-                course_id=self.course_id,
-                role=self.course_role,
-            ).exists()
+            self._assert_user_and_role(self.user.username, self.user.email, self.course_role, self.course_id)
+            # assert CourseStaffRole.objects.filter(
+            #     user=self.user.id,
+            #     course_id=self.course_id,
+            #     role=self.course_role,
+            # ).exists()
 
     def test_add_course_staff_with_new_user(self):
         username = 'pam'
@@ -61,33 +70,65 @@ def test_add_course_staff_with_new_user(self):
         with NamedTemporaryFile() as csv:
             csv = self._write_test_csv(csv, lines)
             call_command(self.command, f'--csv_path={csv.name}')
-            user = User.objects.filter(username=username, email=email)
-            assert user.exists()
-            assert CourseStaffRole.objects.filter(
-                user=user[0].id,
-                course_id=self.course_id,
-                role=self.course_role,
-            ).exists()
+            self._assert_user_and_role(username, email, self.course_role, self.course_id)
+            # user = User.objects.filter(username=username, email=email)
+            # assert user.exists()
+            # assert CourseStaffRole.objects.filter(
+            #     user=user[0].id,
+            #     course_id=self.course_id,
+            #     role=self.course_role,
+            # ).exists()
+
+    def test_add_course_staff_multiple(self):
+        """
+        Assert that the course staff role is correct given multiple lines
+        """
+        username = 'pam'
+        email = '[email protected]'
+        username2 = 'cam'
+        email2 = '[email protected]'
+        lines = [f'{username},{email},{self.course_role},{self.course_id}\n',
+                 f'{username2},{email2},{self.course_role},{self.course_id}\n']
+        with NamedTemporaryFile() as csv:
+            csv = self._write_test_csv(csv, lines)
+            call_command(self.command, f'--csv_path={csv.name}')
+            self._assert_user_and_role(username, email, self.course_role, self.course_id)
+            self._assert_user_and_role(username2, email2, self.course_role, self.course_id)
+            # user = User.objects.filter(username=username, email=email)
+            # assert user.exists()
+            # assert CourseStaffRole.objects.filter(
+            #     user=user[0].id,
+            #     course_id=self.course_id,
+            #     role=self.course_role,
+            # ).exists()
 
     def test_add_course_staff_with_not_default_batch_size(self):
+        """
+        Assert that the number of queries is correct given 2 batches
+        """
         lines = ['pam,[email protected],staff,course-v1:edx+test+f20\n',
                  'sam,[email protected],staff,course-v1:edx+test+f20\n']
         with NamedTemporaryFile() as csv:
             csv = self._write_test_csv(csv, lines)
-            call_command(self.command, f'--csv_path={csv.name}', '--batch_size=1')
-            assert CourseStaffRole.objects.filter(course_id='course-v1:edx+test+f20').count() == 2
+            with self.assertNumQueries(10):
+                call_command(self.command, f'--csv_path={csv.name}', '--batch_size=1')
 
     def test_add_course_staff_with_not_default_batch_delay(self):
-        lines = ['pam,[email protected],staff,course-v1:edx+test+f20\n',
-                 'sam,[email protected],staff,course-v1:edx+test+f20\n']
+        username = 'pam'
+        email = '[email protected]'
+        username2 = 'cam'
+        email2 = '[email protected]'
+        lines = [f'{username},{email},{self.course_role},{self.course_id}\n',
+                 f'{username2},{email2},{self.course_role},{self.course_id}\n']
         with NamedTemporaryFile() as csv:
             csv = self._write_test_csv(csv, lines)
             call_command(self.command, f'--csv_path={csv.name}', '--batch_size=1', '--batch_delay=2')
-            assert CourseStaffRole.objects.filter(course_id='course-v1:edx+test+f20').count() == 2
+            self._assert_user_and_role(username, email, self.course_role, self.course_id)
+            self._assert_user_and_role(username2, email2, self.course_role, self.course_id)
 
     def test_num_queries_correct(self):
         """
-        Expect the number of queries to be 5 + 1 * number of lines
+        Assert the number of queries to be 5 + 1 * number of lines
         - 2 for savepoint/release savepoint, 1 to get existing usernames,
         - 1 to bulk create users, 1 to bulk create course role
         - 1 for each user (to get user)