fix(api): mark attack paths scan as failed when celery task fails (#10065)

api/src/backend/tasks/tests/test_attack_paths_scan.py

@@ -244,9 +244,91 @@ class TestAttackPathsRun:
         failure_args = mock_finish.call_args[0]
         assert failure_args[0] is attack_paths_scan
         assert failure_args[1] == StateChoices.FAILED
-        assert failure_args[2] == {
-            "global_cartography_error": "Cartography failed: ingestion boom"
-        }
+        assert failure_args[2] == {"global_error": "Cartography failed: ingestion boom"}
+
+    def test_run_failure_marks_scan_failed_even_when_drop_database_fails(
+        self, tenants_fixture, providers_fixture, scans_fixture
+    ):
+        tenant = tenants_fixture[0]
+        provider = providers_fixture[0]
+        provider.provider = Provider.ProviderChoices.AWS
+        provider.save()
+        scan = scans_fixture[0]
+        scan.provider = provider
+        scan.save()
+
+        attack_paths_scan = AttackPathsScan.objects.create(
+            tenant_id=tenant.id,
+            provider=provider,
+            scan=scan,
+            state=StateChoices.SCHEDULED,
+        )
+
+        mock_session = MagicMock()
+        session_ctx = MagicMock()
+        session_ctx.__enter__.return_value = mock_session
+        session_ctx.__exit__.return_value = False
+        ingestion_fn = MagicMock(side_effect=RuntimeError("ingestion boom"))
+
+        with (
+            patch(
+                "tasks.jobs.attack_paths.scan.rls_transaction",
+                new=lambda *args, **kwargs: nullcontext(),
+            ),
+            patch(
+                "tasks.jobs.attack_paths.scan.initialize_prowler_provider",
+                return_value=MagicMock(_enabled_regions=["us-east-1"]),
+            ),
+            patch("tasks.jobs.attack_paths.scan.graph_database.get_uri"),
+            patch(
+                "tasks.jobs.attack_paths.scan.graph_database.get_database_name",
+                return_value="db-scan-id",
+            ),
+            patch("tasks.jobs.attack_paths.scan.graph_database.create_database"),
+            patch(
+                "tasks.jobs.attack_paths.scan.graph_database.get_session",
+                return_value=session_ctx,
+            ),
+            patch("tasks.jobs.attack_paths.scan.cartography_create_indexes.run"),
+            patch("tasks.jobs.attack_paths.scan.cartography_analysis.run"),
+            patch("tasks.jobs.attack_paths.scan.findings.create_findings_indexes"),
+            patch("tasks.jobs.attack_paths.scan.internet.analysis"),
+            patch("tasks.jobs.attack_paths.scan.findings.analysis"),
+            patch(
+                "tasks.jobs.attack_paths.scan.db_utils.retrieve_attack_paths_scan",
+                return_value=attack_paths_scan,
+            ),
+            patch("tasks.jobs.attack_paths.scan.db_utils.starting_attack_paths_scan"),
+            patch(
+                "tasks.jobs.attack_paths.scan.db_utils.update_attack_paths_scan_progress"
+            ),
+            patch(
+                "tasks.jobs.attack_paths.scan.db_utils.finish_attack_paths_scan"
+            ) as mock_finish,
+            patch(
+                "tasks.jobs.attack_paths.scan.graph_database.drop_database",
+                side_effect=ConnectionError("neo4j down"),
+            ),
+            patch(
+                "tasks.jobs.attack_paths.scan.get_cartography_ingestion_function",
+                return_value=ingestion_fn,
+            ),
+            patch(
+                "tasks.jobs.attack_paths.scan.utils.call_within_event_loop",
+                side_effect=lambda fn, *a, **kw: fn(*a, **kw),
+            ),
+            patch(
+                "tasks.jobs.attack_paths.scan.utils.stringify_exception",
+                return_value="Cartography failed: ingestion boom",
+            ),
+        ):
+            with pytest.raises(RuntimeError, match="ingestion boom"):
+                attack_paths_run(str(tenant.id), str(scan.id), "task-789")
+
+        failure_args = mock_finish.call_args[0]
+        assert failure_args[0] is attack_paths_scan
+        assert failure_args[1] == StateChoices.FAILED
+        assert failure_args[2] == {"global_error": "Cartography failed: ingestion boom"}
 
     def test_run_returns_early_for_unsupported_provider(self, tenants_fixture):
         tenant = tenants_fixture[0]
@@ -291,6 +373,142 @@ class TestAttackPathsRun:
         mock_retrieve.assert_called_once_with(str(tenant.id), str(scan.id))
 
 
+@pytest.mark.django_db
+class TestFailAttackPathsScan:
+    def test_marks_executing_scan_as_failed(
+        self, tenants_fixture, providers_fixture, scans_fixture
+    ):
+        from tasks.jobs.attack_paths.db_utils import (
+            fail_attack_paths_scan,
+        )
+
+        tenant = tenants_fixture[0]
+        provider = providers_fixture[0]
+        provider.provider = Provider.ProviderChoices.AWS
+        provider.save()
+        scan = scans_fixture[0]
+        scan.provider = provider
+        scan.save()
+
+        attack_paths_scan = AttackPathsScan.objects.create(
+            tenant_id=tenant.id,
+            provider=provider,
+            scan=scan,
+            state=StateChoices.EXECUTING,
+        )
+
+        with (
+            patch(
+                "tasks.jobs.attack_paths.db_utils.retrieve_attack_paths_scan",
+                return_value=attack_paths_scan,
+            ) as mock_retrieve,
+            patch(
+                "tasks.jobs.attack_paths.db_utils.finish_attack_paths_scan"
+            ) as mock_finish,
+        ):
+            fail_attack_paths_scan(str(tenant.id), str(scan.id), "setup exploded")
+
+        mock_retrieve.assert_called_once_with(str(tenant.id), str(scan.id))
+        mock_finish.assert_called_once_with(
+            attack_paths_scan,
+            StateChoices.FAILED,
+            {"global_error": "setup exploded"},
+        )
+
+    def test_skips_already_failed_scan(
+        self, tenants_fixture, providers_fixture, scans_fixture
+    ):
+        from tasks.jobs.attack_paths.db_utils import (
+            fail_attack_paths_scan,
+        )
+
+        tenant = tenants_fixture[0]
+        provider = providers_fixture[0]
+        provider.provider = Provider.ProviderChoices.AWS
+        provider.save()
+        scan = scans_fixture[0]
+        scan.provider = provider
+        scan.save()
+
+        attack_paths_scan = AttackPathsScan.objects.create(
+            tenant_id=tenant.id,
+            provider=provider,
+            scan=scan,
+            state=StateChoices.FAILED,
+        )
+
+        with (
+            patch(
+                "tasks.jobs.attack_paths.db_utils.retrieve_attack_paths_scan",
+                return_value=attack_paths_scan,
+            ),
+            patch(
+                "tasks.jobs.attack_paths.db_utils.finish_attack_paths_scan"
+            ) as mock_finish,
+        ):
+            fail_attack_paths_scan(str(tenant.id), str(scan.id), "setup exploded")
+
+        mock_finish.assert_not_called()
+
+    def test_skips_when_no_scan_found(self, tenants_fixture):
+        from tasks.jobs.attack_paths.db_utils import (
+            fail_attack_paths_scan,
+        )
+
+        tenant = tenants_fixture[0]
+
+        with (
+            patch(
+                "tasks.jobs.attack_paths.db_utils.retrieve_attack_paths_scan",
+                return_value=None,
+            ),
+            patch(
+                "tasks.jobs.attack_paths.db_utils.finish_attack_paths_scan"
+            ) as mock_finish,
+        ):
+            fail_attack_paths_scan(str(tenant.id), "nonexistent", "setup exploded")
+
+        mock_finish.assert_not_called()
+
+
+class TestAttackPathsScanRLSTaskOnFailure:
+    def test_on_failure_delegates_to_fail_attack_paths_scan(self):
+        from tasks.tasks import AttackPathsScanRLSTask
+
+        task = AttackPathsScanRLSTask()
+
+        with patch(
+            "tasks.tasks.attack_paths_db_utils.fail_attack_paths_scan"
+        ) as mock_fail:
+            task.on_failure(
+                exc=RuntimeError("boom"),
+                task_id="task-abc",
+                args=(),
+                kwargs={"tenant_id": "t-1", "scan_id": "s-1"},
+                _einfo=None,
+            )
+
+        mock_fail.assert_called_once_with("t-1", "s-1", "boom")
+
+    def test_on_failure_skips_when_missing_kwargs(self):
+        from tasks.tasks import AttackPathsScanRLSTask
+
+        task = AttackPathsScanRLSTask()
+
+        with patch(
+            "tasks.tasks.attack_paths_db_utils.fail_attack_paths_scan"
+        ) as mock_fail:
+            task.on_failure(
+                exc=RuntimeError("boom"),
+                task_id="task-abc",
+                args=(),
+                kwargs={},
+                _einfo=None,
+            )
+
+        mock_fail.assert_not_called()
+
+
 @pytest.mark.django_db
 class TestAttackPathsFindingsHelpers:
     def test_create_findings_indexes_executes_all_statements(self):