test: fix tests

edx_exams/apps/core/management/commands/bulk_add_course_staff.py

@@ -6,10 +6,7 @@
 from django.core.management.base import BaseCommand
 from django.db import transaction
 
-from edx_exams.apps.core.models import (
-    CourseStaffRole,
-    User,
-)
+from edx_exams.apps.core.models import CourseStaffRole, User
 
 logger = logging.getLogger(__name__)
 

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

@@ -4,9 +4,7 @@
 from django.core.management import call_command
 from django.test import TestCase
 
-from edx_exams.apps.core.models import (
-    CourseStaffRole,
-)
+from edx_exams.apps.core.models import CourseStaffRole
 from edx_exams.apps.core.test_utils.factories import UserFactory
 
 
@@ -43,36 +41,30 @@ def test_empty_csv(self):
             csv = self._write_test_csv(csv, lines)
             with self.assertNumQueries(1):
                 call_command(self.command, f'--csv_path={csv.name}')
-                # todo: figure out how to check success but not via logs
-                # todo: assert number of queries is correct
 
     def test_add_course_staff_with_existing_user(self):
         # todo: create course staff with existing user
         lines = ['amy,[email protected],staff,course-v1:edx+test+f19\n']
         with NamedTemporaryFile() as csv:
             csv = self._write_test_csv(csv, lines)
-            # todo: figure out how to check success but not via logs
             # todo: assert number of queries is correct
-            call_command(self.command, f'--csv_path={csv.name}')
-            assert CourseStaffRole.objects.filter(user=self.user.id).exists()
+            with self.assertNumQueries(2):
+                call_command(self.command, f'--csv_path={csv.name}')
+                assert CourseStaffRole.objects.filter(user=self.user.id).exists()
 
     def test_add_course_staff_with_new_user(self):
-        # todo: create course staff with new user
         lines = ['pam,[email protected],staff,course-v1:edx+test+f20\n']
         with NamedTemporaryFile() as csv:
             csv = self._write_test_csv(csv, lines)
-            # todo: figure out how to check success but not via logs
-            # todo: assert number of queries is correct
-            call_command(self.command, f'--csv_path={csv.name}')
-            assert CourseStaffRole.objects.filter(course_id='course-v1:edx+test+f20').count() == 1
+            with self.assertNumQueries(2):
+                call_command(self.command, f'--csv_path={csv.name}')
+                assert CourseStaffRole.objects.filter(course_id='course-v1:edx+test+f20').count() == 1
 
     def test_add_course_staff_with_batch_size(self):
-        # todo: test with setting batch size
         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)
-            # todo: figure out how to check success but not via logs
             # todo: assert number of queries is correct
             with self.assertNumQueries(2):
                 call_command(self.command, f'--csv_path={csv.name}', '--batch_size=1')