Add Notification to Admins for Ingestor Failure

django_project/core/models/background_task.py

@@ -285,6 +285,16 @@ def task_on_errors(self, exception=None, traceback=None):
                 'progress_text'
             ]
         )
+        if self.task_name in {"ingestor_session", "collector_session"}:
+            session_id = int(self.context_id)
+
+            if self.task_name == "ingestor_session":
+                from gap.tasks.ingestor import notify_ingestor_failure
+                notify_ingestor_failure.delay(session_id, str(exception))
+
+            else:  # "collector_session"
+                from gap.tasks.collector import notify_collector_failure
+                notify_collector_failure.delay(session_id, str(exception))
 
     def task_on_retried(self, reason):
         """Event handler when task is retried by celery.

django_project/core/tests/test_background_task.py

@@ -11,6 +11,7 @@
 from ast import literal_eval as make_tuple
 from django.test import TestCase
 from django.utils import timezone
+from django.conf import settings
 
 from core.models import BackgroundTask, TaskStatus
 from core.celery import (
@@ -30,6 +31,14 @@
 )
 from core.factories import BackgroundTaskF, UserF
 
+from gap.tasks.ingestor import (
+    run_ingestor_session, notify_ingestor_failure
+)
+from gap.tasks.collector import (
+    run_collector_session, notify_collector_failure
+)
+from gap.models.ingestor import IngestorSession, CollectorSession
+
 
 mocked_dt = datetime.datetime(2024, 8, 14, 10, 10, 10, tzinfo=pytz.UTC)
 
@@ -297,3 +306,303 @@ def test_is_task_ignored(self):
         self.assertFalse(
             is_task_ignored("test-new-task")
         )
+
+    @mock.patch("gap.tasks.ingestor.notify_ingestor_failure.delay")
+    def test_task_on_errors_ingestor_session(self, mock_notify):
+        """Test that is triggered when ingestor_session fails."""
+        bg_task = BackgroundTaskF.create(
+            task_name="ingestor_session",
+            context_id="10"
+        )
+        bg_task.task_on_errors(exception="Test ingestor failure")
+        mock_notify.assert_called_once_with(10, "Test ingestor failure")
+
+    @mock.patch("gap.tasks.ingestor.notify_ingestor_failure.delay")
+    def test_task_on_errors_other_tasks(self, mock_notify):
+        """Test do not trigger failure notify."""
+        bg_task = BackgroundTaskF.create(
+            task_name="random_task",
+            context_id="20"
+        )
+        bg_task.task_on_errors(exception="Test failure")
+        mock_notify.assert_not_called()
+
+    @mock.patch("gap.tasks.collector.notify_collector_failure.delay")
+    def test_task_on_errors_collector_session_exception(self, mock_notify):
+        """Test triggered when CollectorSession fails with an error."""
+        bg_task = BackgroundTaskF.create(
+            task_name="collector_session",
+            context_id="25"
+        )
+
+        bg_task.task_on_errors(exception="Test collector failure")
+
+        mock_notify.assert_called_once_with(25, "Test collector failure")
+
+    @mock.patch("gap.tasks.ingestor.notify_ingestor_failure.delay")
+    def test_task_on_errors_ingestor_session_exception(self, mock_notify):
+        """Test triggered when IngestorSession fails with an error."""
+        bg_task = BackgroundTaskF.create(
+            task_name="ingestor_session",
+            context_id="30"
+        )
+
+        bg_task.task_on_errors(exception="Test ingestor failure")
+
+        mock_notify.assert_called_once_with(30, "Test ingestor failure")
+
+    @mock.patch("gap.tasks.ingestor.notify_ingestor_failure.delay")
+    def test_notify_ingestor_failure_session_not_found(self, mock_notify):
+        """Test when an ingestor session does not exist."""
+        with self.assertLogs("gap.tasks.ingestor", level="WARNING") as cm:
+            notify_ingestor_failure(9999, "Session not found")
+
+        self.assertIn("IngestorSession 9999 not found.", cm.output[0])
+
+    @mock.patch("gap.tasks.ingestor.IngestorSession.objects.get")
+    @mock.patch("gap.tasks.ingestor.logger")
+    @mock.patch("core.models.BackgroundTask.objects.filter")
+    def test_notify_ingestor_failure_no_background_task(
+        self, mock_background_task_filter, mock_logger, mock_ingestor_get
+    ):
+        """Test when no BackgroundTask exists for an ingestor session."""
+        # Mock IngestorSession.objects.get to return a dummy session
+        mock_ingestor_get.return_value = mock.Mock()
+
+        # Mock BackgroundTask filter to return None (task does not exist)
+        mock_background_task_filter.return_value.first.return_value = None
+
+        # Call function
+        notify_ingestor_failure(42, "Test failure")
+
+        # Ensure logger warning is logged
+        mock_logger.warning.assert_any_call(
+            "No BackgroundTask found for session 42"
+        )
+
+    @mock.patch("gap.tasks.ingestor.IngestorSession.objects.get")
+    @mock.patch("gap.tasks.ingestor.logger")
+    @mock.patch("django.contrib.auth.get_user_model")
+    def test_notify_ingestor_failure_no_admin_email(
+        self, mock_get_user_model, mock_logger, mock_ingestor_get
+    ):
+        """Test when no admin emails exist."""
+        # Mock IngestorSession.objects.get to return a dummy session
+        mock_ingestor_get.return_value = mock.Mock()
+
+        # Mock user model query to return empty list (no admin emails)
+        mock_user_manager = mock_get_user_model.return_value.objects
+        mock_filtered_users = mock_user_manager.filter.return_value
+        mock_filtered_users.values_list.return_value = []
+
+        # Call function
+        notify_ingestor_failure(42, "Test failure")
+
+        # Verify log message when no admin emails exist
+        mock_logger.warning.assert_any_call(
+            "No admin email found in settings.ADMINS"
+        )
+
+    @mock.patch("gap.tasks.ingestor.notify_ingestor_failure.delay")
+    @mock.patch("gap.models.ingestor.IngestorSession.objects.get")
+    def test_run_ingestor_session_not_found(self, mock_get, mock_notify):
+        """Test triggered when session is not found."""
+        mock_get.side_effect = IngestorSession.DoesNotExist
+
+        run_ingestor_session(9999)
+
+        mock_notify.assert_called_once_with(9999, "Session not found")
+
+    @mock.patch("gap.tasks.ingestor.IngestorSession.objects.get")
+    @mock.patch("gap.tasks.ingestor.logger")
+    @mock.patch("core.models.BackgroundTask.objects.filter")
+    def test_notify_ingestor_failure_with_existing_background_task(
+        self, mock_background_task_filter, mock_logger, mock_ingestor_get
+    ):
+        """Test that BackgroundTask is updated when it exists."""
+        # Mock IngestorSession.objects.get to return a dummy session
+        mock_ingestor_get.return_value = mock.Mock()
+
+        # Create a mock BackgroundTask instance
+        mock_task = mock.Mock()
+        mock_background_task_filter.return_value.first.return_value = mock_task
+
+        # Call function
+        notify_ingestor_failure(42, "Test failure")
+
+        # Ensure status, errors, and last_update were updated
+        self.assertEqual(mock_task.status, TaskStatus.STOPPED)
+        self.assertEqual(mock_task.errors, "Test failure")
+        self.assertTrue(mock_task.last_update)
+
+        # Ensure save() was called once with expected update fields
+        mock_task.save.assert_called_once_with(
+            update_fields=["status", "errors", "last_update"]
+        )
+
+    @mock.patch("gap.tasks.ingestor.send_mail")
+    @mock.patch("gap.tasks.ingestor.IngestorSession.objects.get")
+    @mock.patch("gap.tasks.ingestor.get_user_model")
+    def test_notify_ingestor_failure_with_admin_emails(
+        self, mock_get_user_model, mock_ingestor_get, mock_send_mail
+    ):
+        """Test that email is sent when admin emails exist."""
+        # Mock IngestorSession.objects.get to return a dummy session
+        mock_ingestor_get.return_value = mock.Mock()
+
+        # Mock user model query to return a list of admin emails
+        mock_user_manager = mock_get_user_model.return_value.objects
+        mock_filtered_users = mock_user_manager.filter.return_value
+        mock_filtered_users.values_list.return_value = ["[email protected]"]
+
+        # Call function
+        notify_ingestor_failure(42, "Test failure")
+
+        # Ensure send_mail() was called with correct parameters
+        mock_send_mail.assert_called_once_with(
+            subject="Ingestor Failure Alert",
+            message=(
+                "Ingestor Session 42 has failed.\n\n"
+                "Error: Test failure\n\n"
+                "Please check the logs for more details."
+            ),
+            from_email=settings.DEFAULT_FROM_EMAIL,
+            recipient_list=[["[email protected]"]],
+            fail_silently=False,
+        )
+
+    @mock.patch("gap.tasks.ingestor.IngestorSession.objects.get")
+    @mock.patch("gap.tasks.ingestor.get_user_model")
+    @mock.patch("gap.tasks.ingestor.send_mail")
+    def test_notify_ingestor_failure_return_value(
+        self, mock_send_mail, mock_get_user_model, mock_ingestor_get
+    ):
+        """Test that notify_ingestor_failure returns the expected string."""
+        # Mock IngestorSession.objects.get
+        mock_session = mock.Mock()
+        mock_ingestor_get.return_value = mock_session
+
+        # Mock user model query to return an admin email
+        mock_user_manager = mock_get_user_model.return_value.objects
+        mock_filtered_users = mock_user_manager.filter.return_value
+        mock_filtered_users.values_list.return_value = ["[email protected]"]
+
+        # Call function and store return value
+        result = notify_ingestor_failure(42, "Test failure")
+
+        # Expected return message
+        expected_msg = (
+            "Logged ingestor failure for session 42 and notified admins"
+        )
+
+        # Assert function returned the expected message
+        self.assertEqual(result, expected_msg)
+
+    @mock.patch("gap.models.ingestor.CollectorSession.objects.get")
+    @mock.patch("gap.tasks.collector.notify_collector_failure.delay")
+    def test_run_collector_session_not_found(self, mock_notify, mock_get):
+        """Test triggered when collector session is not found."""
+        mock_get.side_effect = CollectorSession.DoesNotExist
+
+        run_collector_session(9999)
+
+        mock_notify.assert_called_once_with(
+            9999, "Collector session not found"
+        )
+
+    @mock.patch("gap.tasks.collector.notify_collector_failure.delay")
+    def test_task_on_errors_collector_session(self, mock_notify):
+        """Test triggered when collector_session fails."""
+        bg_task = BackgroundTask.objects.create(
+            task_name="collector_session",
+            context_id="15"
+        )
+        bg_task.task_on_errors(exception="Test collector failure")
+
+        mock_notify.assert_called_once_with(15, "Test collector failure")
+
+    @mock.patch("gap.tasks.collector.CollectorSession.objects.get")
+    @mock.patch("gap.tasks.collector.logger")
+    @mock.patch("core.models.BackgroundTask.objects.filter")
+    def test_notify_collector_failure_no_background_task(
+        self, mock_background_task_filter, mock_logger, mock_collector_get
+    ):
+        """Test when no BackgroundTask exists for a collector session."""
+        mock_collector_get.return_value = mock.Mock()
+
+        mock_background_task_filter.return_value.first.return_value = None
+
+        notify_collector_failure(42, "Test failure")
+        print(mock_logger.mock_calls)
+
+        mock_logger.warning.assert_any_call(
+            "No BackgroundTask found for collector session 42"
+        )
+
+    @mock.patch("gap.tasks.collector.CollectorSession.objects.get")
+    @mock.patch("gap.tasks.collector.logger")
+    @mock.patch("django.contrib.auth.get_user_model")
+    def test_notify_collector_failure_no_admin_email(
+        self, mock_get_user_model, mock_logger, mock_collector_get
+    ):
+        """Test when no admin emails exist for collector failure."""
+        mock_collector_get.return_value = mock.Mock()
+
+        mock_user_manager = mock_get_user_model.return_value.objects
+        mock_filtered_users = mock_user_manager.filter.return_value
+        mock_filtered_users.values_list.return_value = []
+
+        notify_collector_failure(42, "Test failure")
+
+        mock_logger.warning.assert_any_call(
+            "No admin email found in settings.ADMINS"
+        )
+
+    @mock.patch("gap.tasks.collector.send_mail")
+    @mock.patch("gap.tasks.collector.CollectorSession.objects.get")
+    @mock.patch("gap.tasks.collector.get_user_model")
+    def test_notify_collector_failure_with_admin_emails(
+        self, mock_get_user_model, mock_collector_get, mock_send_mail
+    ):
+        """Test that email is sent for collector failure."""
+        mock_collector_get.return_value = mock.Mock()
+
+        mock_user_manager = mock_get_user_model.return_value.objects
+        mock_filtered_users = mock_user_manager.filter.return_value
+        mock_filtered_users.values_list.return_value = [["[email protected]"]]
+
+        notify_collector_failure(42, "Test failure")
+
+        mock_send_mail.assert_called_once_with(
+            subject="Collector Failure Alert",
+            message=(
+                "Collector Session 42 has failed.\n\n"
+                "Error: Test failure\n\n"
+                "Please check the logs for more details."
+            ),
+            from_email=settings.DEFAULT_FROM_EMAIL,
+            recipient_list=[["[email protected]"]],
+            fail_silently=False,
+        )
+
+    @mock.patch("gap.tasks.collector.CollectorSession.objects.get")
+    @mock.patch("gap.tasks.collector.get_user_model")
+    @mock.patch("gap.tasks.collector.send_mail")
+    def test_notify_collector_failure_return_value(
+        self, mock_send_mail, mock_get_user_model, mock_collector_get
+    ):
+        """Test that notify_collector_failure returns the expected string."""
+        mock_session = mock.Mock()
+        mock_collector_get.return_value = mock_session
+
+        mock_user_manager = mock_get_user_model.return_value.objects
+        mock_filtered_users = mock_user_manager.filter.return_value
+        mock_filtered_users.values_list.return_value = [["[email protected]"]]
+
+        result = notify_collector_failure(42, "Test failure")
+
+        expected_msg = (
+            "Logged collector 42 failed. Admins notified."
+        )
+
+        self.assertEqual(result, expected_msg)