From 10f22834ee114edc732a35a20189318dc49190ba Mon Sep 17 00:00:00 2001
From: Trenton H <797416+stumpylog@users.noreply.github.com>
Date: Thu, 9 Nov 2023 11:21:45 -0800
Subject: [PATCH] Use a single process in testing, not the Pool

---
 .../management/commands/document_archiver.py  | 24 ++++++++++++-------
 .../commands/document_fuzzy_match.py          |  6 ++---
 .../commands/document_thumbnails.py           | 20 +++++++++-------
 src/documents/tests/test_management.py        |  2 +-
 .../tests/test_management_thumbnails.py       |  4 ++--
 5 files changed, 34 insertions(+), 22 deletions(-)

diff --git a/src/documents/management/commands/document_archiver.py b/src/documents/management/commands/document_archiver.py
index 40714e866..96b3b50ab 100644
--- a/src/documents/management/commands/document_archiver.py
+++ b/src/documents/management/commands/document_archiver.py
@@ -75,13 +75,21 @@ class Command(MultiProcessMixin, ProgressBarMixin, BaseCommand):
 
         try:
             logging.getLogger().handlers[0].level = logging.ERROR
-            with multiprocessing.Pool(self.process_count) as pool:
-                list(
-                    tqdm.tqdm(
-                        pool.imap_unordered(update_document_archive_file, document_ids),
-                        total=len(document_ids),
-                        disable=self.no_progress_bar,
-                    ),
-                )
+
+            if self.process_count == 1:
+                for doc_id in document_ids:
+                    update_document_archive_file(doc_id)
+            else:  # pragma: no cover
+                with multiprocessing.Pool(self.process_count) as pool:
+                    list(
+                        tqdm.tqdm(
+                            pool.imap_unordered(
+                                update_document_archive_file,
+                                document_ids,
+                            ),
+                            total=len(document_ids),
+                            disable=self.no_progress_bar,
+                        ),
+                    )
         except KeyboardInterrupt:
             self.stdout.write(self.style.NOTICE("Aborting..."))
diff --git a/src/documents/management/commands/document_fuzzy_match.py b/src/documents/management/commands/document_fuzzy_match.py
index 63c640bee..597a9d2c1 100644
--- a/src/documents/management/commands/document_fuzzy_match.py
+++ b/src/documents/management/commands/document_fuzzy_match.py
@@ -93,12 +93,12 @@ class Command(MultiProcessMixin, ProgressBarMixin, BaseCommand):
                 work_pkgs.append(_WorkPackage(first_doc, second_doc))
 
         # Don't spin up a pool of 1 process
-        if options["processes"] == 1:
+        if self.process_count == 1:
             results = []
             for work in tqdm.tqdm(work_pkgs, disable=self.no_progress_bar):
                 results.append(_process_and_match(work))
-        else:
-            with multiprocessing.Pool(processes=options["processes"]) as pool:
+        else:  # pragma: no cover
+            with multiprocessing.Pool(processes=self.process_count) as pool:
                 results = list(
                     tqdm.tqdm(
                         pool.imap_unordered(_process_and_match, work_pkgs),
diff --git a/src/documents/management/commands/document_thumbnails.py b/src/documents/management/commands/document_thumbnails.py
index fdc3d319e..ecd265102 100644
--- a/src/documents/management/commands/document_thumbnails.py
+++ b/src/documents/management/commands/document_thumbnails.py
@@ -70,11 +70,15 @@ class Command(MultiProcessMixin, ProgressBarMixin, BaseCommand):
         # with postgres.
         db.connections.close_all()
 
-        with multiprocessing.Pool(processes=self.process_count) as pool:
-            list(
-                tqdm.tqdm(
-                    pool.imap_unordered(_process_document, ids),
-                    total=len(ids),
-                    disable=self.no_progress_bar,
-                ),
-            )
+        if self.process_count == 1:
+            for doc_id in ids:
+                _process_document(doc_id)
+        else:  # pragma: no cover
+            with multiprocessing.Pool(processes=self.process_count) as pool:
+                list(
+                    tqdm.tqdm(
+                        pool.imap_unordered(_process_document, ids),
+                        total=len(ids),
+                        disable=self.no_progress_bar,
+                    ),
+                )
diff --git a/src/documents/tests/test_management.py b/src/documents/tests/test_management.py
index 1115325ca..d1efe27d4 100644
--- a/src/documents/tests/test_management.py
+++ b/src/documents/tests/test_management.py
@@ -36,7 +36,7 @@ class TestArchiver(DirectoriesMixin, FileSystemAssertsMixin, TestCase):
             os.path.join(self.dirs.originals_dir, f"{doc.id:07}.pdf"),
         )
 
-        call_command("document_archiver")
+        call_command("document_archiver", "--processes", "1")
 
     def test_handle_document(self):
         doc = self.make_models()
diff --git a/src/documents/tests/test_management_thumbnails.py b/src/documents/tests/test_management_thumbnails.py
index 9f3ff63c5..4056b65fe 100644
--- a/src/documents/tests/test_management_thumbnails.py
+++ b/src/documents/tests/test_management_thumbnails.py
@@ -83,13 +83,13 @@ class TestMakeThumbnails(DirectoriesMixin, FileSystemAssertsMixin, TestCase):
     def test_command(self):
         self.assertIsNotFile(self.d1.thumbnail_path)
         self.assertIsNotFile(self.d2.thumbnail_path)
-        call_command("document_thumbnails")
+        call_command("document_thumbnails", "--processes", "1")
         self.assertIsFile(self.d1.thumbnail_path)
         self.assertIsFile(self.d2.thumbnail_path)
 
     def test_command_documentid(self):
         self.assertIsNotFile(self.d1.thumbnail_path)
         self.assertIsNotFile(self.d2.thumbnail_path)
-        call_command("document_thumbnails", "-d", f"{self.d1.id}")
+        call_command("document_thumbnails", "--processes", "1", "-d", f"{self.d1.id}")
         self.assertIsFile(self.d1.thumbnail_path)
         self.assertIsNotFile(self.d2.thumbnail_path)