From 8391bf6cb7d1b34fff2f06b8d6d98d5f255f91c4 Mon Sep 17 00:00:00 2001
From: Matthew Elwell <matthew.elwell@flagsmith.com>
Date: Tue, 9 Apr 2024 11:38:53 +0100
Subject: [PATCH] Catch all exceptions

---
 api/task_processor/threads.py                   |  5 ++---
 .../test_unit_task_processor_threads.py         | 17 ++++++++++++-----
 2 files changed, 14 insertions(+), 8 deletions(-)

diff --git a/api/task_processor/threads.py b/api/task_processor/threads.py
index 0cc995274cce..a3838a3310ca 100644
--- a/api/task_processor/threads.py
+++ b/api/task_processor/threads.py
@@ -3,7 +3,6 @@
 import traceback
 from threading import Thread
 
-from django.db import DatabaseError
 from django.utils import timezone
 
 from task_processor.processor import run_recurring_tasks, run_tasks
@@ -36,14 +35,14 @@ def run_iteration(self) -> None:
         try:
             run_tasks(self.queue_pop_size)
             run_recurring_tasks(self.queue_pop_size)
-        except DatabaseError as e:
+        except Exception as e:
             # To prevent task threads from dying if they get an error retrieving the tasks from the
             # database this will allow the thread to continue trying to retrieve tasks if it can
             # successfully re-establish a connection to the database.
             # TODO: is this also what is causing tasks to get stuck as locked? Can we unlock
             #  tasks here?
 
-            logger.error("Received database error retrieving tasks: %s.", e)
+            logger.error("Received error retrieving tasks: %s.", e)
             logger.debug(traceback.format_exc())
 
     def stop(self):
diff --git a/api/tests/unit/task_processor/test_unit_task_processor_threads.py b/api/tests/unit/task_processor/test_unit_task_processor_threads.py
index 391c3fd14325..2957489c8448 100644
--- a/api/tests/unit/task_processor/test_unit_task_processor_threads.py
+++ b/api/tests/unit/task_processor/test_unit_task_processor_threads.py
@@ -1,25 +1,32 @@
 import logging
+from typing import Type
 
+import pytest
 from django.db import DatabaseError
-from pytest_django.fixtures import SettingsWrapper
 from pytest_mock import MockerFixture
 
 from task_processor.threads import TaskRunner
 from tests.unit.task_processor.conftest import GetTaskProcessorCaplog
 
 
-def test_task_runner_is_resilient_to_database_errors(
+@pytest.mark.parametrize(
+    "exception_class, exception_message",
+    [(DatabaseError, "Database error"), (Exception, "Generic error")],
+)
+def test_task_runner_is_resilient_to_errors(
     db: None,
     mocker: MockerFixture,
     get_task_processor_caplog: GetTaskProcessorCaplog,
-    settings: SettingsWrapper,
+    exception_class: Type[Exception],
+    exception_message: str,
 ) -> None:
     # Given
     caplog = get_task_processor_caplog(logging.DEBUG)
 
     task_runner = TaskRunner()
     mocker.patch(
-        "task_processor.threads.run_tasks", side_effect=DatabaseError("Database error")
+        "task_processor.threads.run_tasks",
+        side_effect=exception_class(exception_message),
     )
 
     # When
@@ -31,7 +38,7 @@ def test_task_runner_is_resilient_to_database_errors(
     assert caplog.records[0].levelno == logging.ERROR
     assert (
         caplog.records[0].message
-        == "Received database error retrieving tasks: Database error."
+        == f"Received error retrieving tasks: {exception_message}."
     )
 
     assert caplog.records[1].levelno == logging.DEBUG