add dry_run option to volunteer import script

Update the script to run all changes inside a transaction and if the
dry_run option is passed then rollback all changes.
This lets you see any issues with running out of assignments without
having made any changes to the database.

Author: struan

crowdsourcer/management/commands/import_volunteers.py

@@ -4,6 +4,7 @@
 from django.conf import settings
 from django.contrib.auth.models import User
 from django.core.management.base import BaseCommand
+from django.db import transaction
 from django.db.models import Q
 
 import pandas as pd
@@ -113,6 +114,12 @@ def add_arguments(self, parser):
             "--make_assignments", action="store_true", help="assign councils to users"
         )
 
+        parser.add_argument(
+            "--dry_run",
+            action="store_true",
+            help="run everything and then undo changes. Helpful to check if assigment weighting is good",
+        )
+
     def get_df(self, filename):
         df = pd.read_csv(
             filename,
@@ -149,32 +156,7 @@ def get_assignment_count(self, user_type):
         num_councils = self.num_council_map.get(user_type)
         return num_councils
 
-    def handle(
-        self,
-        quiet: bool = False,
-        file: str = None,
-        session: str = None,
-        col_names: str = None,
-        response_type: str = None,
-        assignment_map: str = None,
-        authority_map: str = None,
-        *args,
-        **options,
-    ):
-        if file is None:
-            file = self.volunteer_file
-
-        self.set_cols(col_names)
-        self.set_assignment_map(assignment_map)
-        self.set_authority_map(authority_map)
-
-        df = self.get_df(file)
-
-        if response_type is None:
-            response_type = self.response_type
-
-        session = MarkingSession.objects.get(label=session)
-        rt = ResponseType.objects.get(type=response_type)
+    def add_users_and_assignments(self, df, response_type, session, rt, options):
 
         bad_councils = []
         for index, row in df.iterrows():
@@ -300,28 +282,70 @@ def handle(
                         response_type=rt,
                     )
 
-        council_count = PublicAuthority.objects.filter(
-            marking_session=session, do_not_mark=False
-        ).count()
-        for section in Section.objects.filter(marking_session=session).all():
-            assigned = Assigned.objects.filter(section=section).count()
-            if assigned != council_count:
-                self.stdout.write(
-                    f"{RED}Not all councils assigned for {section.title} ({assigned}/{council_count}){NOBOLD}"
-                )
-            else:
-                self.stdout.write(f"{GREEN}All councils and sections assigned{NOBOLD}")
-
-        volunteer_count = User.objects.filter(
-            marker__marking_session=session, marker__response_type=rt
-        ).count()
-        assigned_count = (
-            Assigned.objects.filter(user__is_superuser=False)
-            .distinct("user_id")
-            .count()
-        )
+        return bad_councils
+
+    def handle(
+        self,
+        quiet: bool = False,
+        file: str = None,
+        session: str = None,
+        col_names: str = None,
+        response_type: str = None,
+        assignment_map: str = None,
+        authority_map: str = None,
+        *args,
+        **options,
+    ):
+        if file is None:
+            file = self.volunteer_file
+
+        self.set_cols(col_names)
+        self.set_assignment_map(assignment_map)
+        self.set_authority_map(authority_map)
+
+        df = self.get_df(file)
+
+        if response_type is None:
+            response_type = self.response_type
+
+        session = MarkingSession.objects.get(label=session)
+        rt = ResponseType.objects.get(type=response_type)
+
+        with transaction.atomic():
+            bad_councils = self.add_users_and_assignments(
+                df, response_type, session, rt, options
+            )
+
+            council_count = PublicAuthority.objects.filter(
+                marking_session=session, do_not_mark=False
+            ).count()
+            for section in Section.objects.filter(marking_session=session).all():
+                assigned = Assigned.objects.filter(section=section).count()
+                if assigned != council_count:
+                    self.stdout.write(
+                        f"{RED}Not all councils assigned for {section.title} ({assigned}/{council_count}){NOBOLD}"
+                    )
+                else:
+                    self.stdout.write(
+                        f"{GREEN}All councils and sections assigned{NOBOLD}"
+                    )
+
+            volunteer_count = User.objects.filter(
+                marker__marking_session=session, marker__response_type=rt
+            ).count()
+            assigned_count = (
+                Assigned.objects.filter(user__is_superuser=False)
+                .distinct("user_id")
+                .count()
+            )
+
+            self.stdout.write(
+                f"{assigned_count}/{volunteer_count} users assigned marking"
+            )
+
+            if options["dry_run"]:
+                transaction.set_rollback(True)
 
-        self.stdout.write(f"{assigned_count}/{volunteer_count} users assigned marking")
         if not options["add_users"]:
             self.stdout.write(
                 f"{YELLOW}Dry run, no users added, call with --add_users to add users{NOBOLD}"
@@ -331,6 +355,11 @@ def handle(
                 f"{YELLOW}Dry run, no assignments made, call with --make_assignments to make them{NOBOLD}"
             )
 
+        if options["dry_run"]:
+            self.stdout.write(
+                f"{YELLOW}Dry run, no changes made, call without --dry_run to make changes{NOBOLD}"
+            )
+
         if len(bad_councils):
             self.stdout.write("Bad councils are:")
             for c in bad_councils:

crowdsourcer/tests/test_import_volunteers_script.py

@@ -119,6 +119,24 @@ def test_basic_run(self):
         self.assertTrue("Aberdeen City Council" in councils)
         self.assertTrue("Adur District Council" in councils)
 
+    def test_dry_run(self):
+        data_file = pathlib.Path(__file__).parent.resolve() / "data" / "volunteers.csv"
+
+        self.assertEquals(User.objects.count(), 0)
+        self.assertEquals(Marker.objects.count(), 0)
+        self.assertEquals(Assigned.objects.count(), 0)
+        self.call_command(
+            "import_volunteers",
+            session="Default",
+            file=data_file,
+            add_users=True,
+            make_assignments=True,
+            dry_run=True,
+        )
+        self.assertEquals(User.objects.count(), 0)
+        self.assertEquals(Marker.objects.count(), 0)
+        self.assertEquals(Assigned.objects.count(), 0)
+
     def test_multi_councils(self):
         data_file = (
             pathlib.Path(__file__).parent.resolve()