import uuid
from unittest.mock import MagicMock, patch

import pytest
from tasks.tasks import (
    _perform_scan_complete_tasks,
    check_integrations_task,
    generate_outputs_task,
    s3_integration_task,
    security_hub_integration_task,
)

from api.models import Integration


# TODO Move this to outputs/reports jobs
@pytest.mark.django_db
class TestGenerateOutputs:
    def setup_method(self):
        self.scan_id = str(uuid.uuid4())
        self.provider_id = str(uuid.uuid4())
        self.tenant_id = str(uuid.uuid4())

    def test_no_findings_returns_early(self):
        with patch("tasks.tasks.ScanSummary.objects.filter") as mock_filter:
            mock_filter.return_value.exists.return_value = False

            result = generate_outputs_task(
                scan_id=self.scan_id,
                provider_id=self.provider_id,
                tenant_id=self.tenant_id,
            )

            assert result == {"upload": False}
            mock_filter.assert_called_once_with(scan_id=self.scan_id)

    @patch("tasks.tasks._upload_to_s3")
    @patch("tasks.tasks._compress_output_files")
    @patch("tasks.tasks.get_compliance_frameworks")
    @patch("tasks.tasks.Compliance.get_bulk")
    @patch("tasks.tasks.initialize_prowler_provider")
    @patch("tasks.tasks.Provider.objects.get")
    @patch("tasks.tasks.ScanSummary.objects.filter")
    @patch("tasks.tasks.Finding.all_objects.filter")
    def test_generate_outputs_happy_path(
        self,
        mock_finding_filter,
        mock_scan_summary_filter,
        mock_provider_get,
        mock_initialize_provider,
        mock_compliance_get_bulk,
        mock_get_available_frameworks,
        mock_compress,
        mock_upload,
    ):
        mock_scan_summary_filter.return_value.exists.return_value = True

        mock_provider = MagicMock()
        mock_provider.uid = "provider-uid"
        mock_provider.provider = "aws"
        mock_provider_get.return_value = mock_provider

        prowler_provider = MagicMock()
        mock_initialize_provider.return_value = prowler_provider

        mock_compliance_get_bulk.return_value = {"cis": MagicMock()}
        mock_get_available_frameworks.return_value = ["cis"]

        dummy_finding = MagicMock(uid="f1")
        mock_finding_filter.return_value.order_by.return_value.iterator.return_value = [
            [dummy_finding],
            True,
        ]

        mock_transformed_stats = {"some": "stats"}
        with (
            patch(
                "tasks.tasks.FindingOutput._transform_findings_stats",
                return_value=mock_transformed_stats,
            ),
            patch(
                "tasks.tasks.FindingOutput.transform_api_finding",
                return_value={"transformed": "f1"},
            ),
            patch(
                "tasks.tasks.OUTPUT_FORMATS_MAPPING",
                {
                    "json": {
                        "class": MagicMock(name="JSONWriter"),
                        "suffix": ".json",
                        "kwargs": {},
                    }
                },
            ),
            patch(
                "tasks.tasks.COMPLIANCE_CLASS_MAP",
                {"aws": [(lambda x: True, MagicMock(name="CSVCompliance"))]},
            ),
            patch(
                "tasks.tasks._generate_output_directory",
                return_value=("out-dir", "comp-dir"),
            ),
            patch("tasks.tasks.Scan.all_objects.filter") as mock_scan_update,
            patch("tasks.tasks.rmtree"),
        ):
            mock_compress.return_value = "/tmp/zipped.zip"
            mock_upload.return_value = "s3://bucket/zipped.zip"

            result = generate_outputs_task(
                scan_id=self.scan_id,
                provider_id=self.provider_id,
                tenant_id=self.tenant_id,
            )

            assert result == {"upload": True}
            mock_scan_update.return_value.update.assert_called_once_with(
                output_location="s3://bucket/zipped.zip"
            )

    def test_generate_outputs_fails_upload(self):
        with (
            patch("tasks.tasks.ScanSummary.objects.filter") as mock_filter,
            patch("tasks.tasks.Provider.objects.get"),
            patch("tasks.tasks.initialize_prowler_provider"),
            patch("tasks.tasks.Compliance.get_bulk"),
            patch("tasks.tasks.get_compliance_frameworks"),
            patch("tasks.tasks.Finding.all_objects.filter") as mock_findings,
            patch(
                "tasks.tasks._generate_output_directory", return_value=("out", "comp")
            ),
            patch("tasks.tasks.FindingOutput._transform_findings_stats"),
            patch("tasks.tasks.FindingOutput.transform_api_finding"),
            patch(
                "tasks.tasks.OUTPUT_FORMATS_MAPPING",
                {
                    "json": {
                        "class": MagicMock(name="Writer"),
                        "suffix": ".json",
                        "kwargs": {},
                    }
                },
            ),
            patch(
                "tasks.tasks.COMPLIANCE_CLASS_MAP",
                {"aws": [(lambda x: True, MagicMock())]},
            ),
            patch("tasks.tasks._compress_output_files", return_value="/tmp/compressed"),
            patch("tasks.tasks._upload_to_s3", return_value=None),
            patch("tasks.tasks.Scan.all_objects.filter") as mock_scan_update,
            patch("tasks.tasks.rmtree"),
        ):
            mock_filter.return_value.exists.return_value = True
            mock_findings.return_value.order_by.return_value.iterator.return_value = [
                [MagicMock()],
                True,
            ]

            result = generate_outputs_task(
                scan_id="scan",
                provider_id=self.provider_id,
                tenant_id=self.tenant_id,
            )

            assert result == {"upload": False}
            mock_scan_update.return_value.update.assert_called_once()

    def test_generate_outputs_triggers_html_extra_update(self):
        mock_finding_output = MagicMock()
        mock_finding_output.compliance = {"cis": ["requirement-1", "requirement-2"]}

        with (
            patch("tasks.tasks.ScanSummary.objects.filter") as mock_filter,
            patch("tasks.tasks.Provider.objects.get"),
            patch("tasks.tasks.initialize_prowler_provider"),
            patch("tasks.tasks.Compliance.get_bulk", return_value={"cis": MagicMock()}),
            patch("tasks.tasks.get_compliance_frameworks", return_value=["cis"]),
            patch("tasks.tasks.Finding.all_objects.filter") as mock_findings,
            patch(
                "tasks.tasks._generate_output_directory", return_value=("out", "comp")
            ),
            patch(
                "tasks.tasks.FindingOutput._transform_findings_stats",
                return_value={"some": "stats"},
            ),
            patch(
                "tasks.tasks.FindingOutput.transform_api_finding",
                return_value=mock_finding_output,
            ),
            patch("tasks.tasks._compress_output_files", return_value="/tmp/compressed"),
            patch("tasks.tasks._upload_to_s3", return_value="s3://bucket/f.zip"),
            patch("tasks.tasks.Scan.all_objects.filter"),
            patch("tasks.tasks.rmtree"),
        ):
            mock_filter.return_value.exists.return_value = True
            mock_findings.return_value.order_by.return_value.iterator.return_value = [
                [MagicMock()],
                True,
            ]

            html_writer_mock = MagicMock()
            with (
                patch(
                    "tasks.tasks.OUTPUT_FORMATS_MAPPING",
                    {
                        "html": {
                            "class": lambda *args, **kwargs: html_writer_mock,
                            "suffix": ".html",
                            "kwargs": {},
                        }
                    },
                ),
                patch(
                    "tasks.tasks.COMPLIANCE_CLASS_MAP",
                    {"aws": [(lambda x: True, MagicMock())]},
                ),
            ):
                generate_outputs_task(
                    scan_id=self.scan_id,
                    provider_id=self.provider_id,
                    tenant_id=self.tenant_id,
                )
                html_writer_mock.batch_write_data_to_file.assert_called_once()

    def test_transform_called_only_on_second_batch(self):
        raw1 = MagicMock()
        raw2 = MagicMock()

        tf1 = MagicMock()
        tf1.compliance = {}
        tf2 = MagicMock()
        tf2.compliance = {}

        writer_instances = []

        class TrackingWriter:
            def __init__(self, findings, file_path, file_extension, from_cli):
                self.transform_called = 0
                self.batch_write_data_to_file = MagicMock()
                self._data = []
                self.close_file = False
                writer_instances.append(self)

            def transform(self, fos):
                self.transform_called += 1

        with (
            patch("tasks.tasks.ScanSummary.objects.filter") as mock_summary,
            patch("tasks.tasks.Provider.objects.get"),
            patch("tasks.tasks.initialize_prowler_provider"),
            patch("tasks.tasks.Compliance.get_bulk"),
            patch("tasks.tasks.get_compliance_frameworks", return_value=[]),
            patch("tasks.tasks.FindingOutput._transform_findings_stats"),
            patch(
                "tasks.tasks.FindingOutput.transform_api_finding",
                side_effect=[tf1, tf2],
            ),
            patch(
                "tasks.tasks._generate_output_directory",
                return_value=("outdir", "compdir"),
            ),
            patch("tasks.tasks._compress_output_files", return_value="outdir.zip"),
            patch("tasks.tasks._upload_to_s3", return_value="s3://bucket/outdir.zip"),
            patch("tasks.tasks.Scan.all_objects.filter"),
            patch("tasks.tasks.rmtree"),
            patch(
                "tasks.tasks.batched",
                return_value=[
                    ([raw1], False),
                    ([raw2], True),
                ],
            ),
        ):
            mock_summary.return_value.exists.return_value = True

            with patch(
                "tasks.tasks.OUTPUT_FORMATS_MAPPING",
                {
                    "json": {
                        "class": TrackingWriter,
                        "suffix": ".json",
                        "kwargs": {},
                    }
                },
            ):
                result = generate_outputs_task(
                    scan_id=self.scan_id,
                    provider_id=self.provider_id,
                    tenant_id=self.tenant_id,
                )

        assert result == {"upload": True}
        assert len(writer_instances) == 1
        writer = writer_instances[0]
        assert writer.transform_called == 1

    def test_compliance_transform_called_on_second_batch(self):
        raw1 = MagicMock()
        raw2 = MagicMock()
        compliance_obj = MagicMock()
        writer_instances = []

        class TrackingComplianceWriter:
            def __init__(self, *args, **kwargs):
                self.transform_calls = []
                self._data = []
                writer_instances.append(self)

            def transform(self, fos, comp_obj, name):
                self.transform_calls.append((fos, comp_obj, name))

            def batch_write_data_to_file(self):
                pass

        two_batches = [
            ([raw1], False),
            ([raw2], True),
        ]

        with (
            patch("tasks.tasks.ScanSummary.objects.filter") as mock_summary,
            patch(
                "tasks.tasks.Provider.objects.get",
                return_value=MagicMock(uid="UID", provider="aws"),
            ),
            patch("tasks.tasks.initialize_prowler_provider"),
            patch(
                "tasks.tasks.Compliance.get_bulk", return_value={"cis": compliance_obj}
            ),
            patch("tasks.tasks.get_compliance_frameworks", return_value=["cis"]),
            patch(
                "tasks.tasks._generate_output_directory",
                return_value=("outdir", "compdir"),
            ),
            patch("tasks.tasks.FindingOutput._transform_findings_stats"),
            patch(
                "tasks.tasks.FindingOutput.transform_api_finding",
                side_effect=lambda f, prov: f,
            ),
            patch("tasks.tasks._compress_output_files", return_value="outdir.zip"),
            patch("tasks.tasks._upload_to_s3", return_value="s3://bucket/outdir.zip"),
            patch(
                "tasks.tasks.Scan.all_objects.filter",
                return_value=MagicMock(update=lambda **kw: None),
            ),
            patch("tasks.tasks.batched", return_value=two_batches),
            patch("tasks.tasks.OUTPUT_FORMATS_MAPPING", {}),
            patch("tasks.tasks.rmtree"),
            patch(
                "tasks.tasks.COMPLIANCE_CLASS_MAP",
                {"aws": [(lambda name: True, TrackingComplianceWriter)]},
            ),
        ):
            mock_summary.return_value.exists.return_value = True

            result = generate_outputs_task(
                scan_id=self.scan_id,
                provider_id=self.provider_id,
                tenant_id=self.tenant_id,
            )

        assert len(writer_instances) == 1
        writer = writer_instances[0]
        assert writer.transform_calls == [([raw2], compliance_obj, "cis")]
        assert result == {"upload": True}

    # TODO: We need to add a periodic task to delete old output files
    def test_generate_outputs_logs_rmtree_exception(self, caplog):
        mock_finding_output = MagicMock()
        mock_finding_output.compliance = {"cis": ["requirement-1", "requirement-2"]}

        with (
            patch("tasks.tasks.ScanSummary.objects.filter") as mock_filter,
            patch("tasks.tasks.Provider.objects.get"),
            patch("tasks.tasks.initialize_prowler_provider"),
            patch("tasks.tasks.Compliance.get_bulk", return_value={"cis": MagicMock()}),
            patch("tasks.tasks.get_compliance_frameworks", return_value=["cis"]),
            patch("tasks.tasks.Finding.all_objects.filter") as mock_findings,
            patch(
                "tasks.tasks._generate_output_directory", return_value=("out", "comp")
            ),
            patch(
                "tasks.tasks.FindingOutput._transform_findings_stats",
                return_value={"some": "stats"},
            ),
            patch(
                "tasks.tasks.FindingOutput.transform_api_finding",
                return_value=mock_finding_output,
            ),
            patch("tasks.tasks._compress_output_files", return_value="/tmp/compressed"),
            patch("tasks.tasks._upload_to_s3", return_value="s3://bucket/file.zip"),
            patch("tasks.tasks.Scan.all_objects.filter"),
            patch("tasks.tasks.rmtree", side_effect=Exception("Test deletion error")),
        ):
            mock_filter.return_value.exists.return_value = True
            mock_findings.return_value.order_by.return_value.iterator.return_value = [
                [MagicMock()],
                True,
            ]

            with (
                patch(
                    "tasks.tasks.OUTPUT_FORMATS_MAPPING",
                    {
                        "json": {
                            "class": lambda *args, **kwargs: MagicMock(),
                            "suffix": ".json",
                            "kwargs": {},
                        }
                    },
                ),
                patch(
                    "tasks.tasks.COMPLIANCE_CLASS_MAP",
                    {"aws": [(lambda x: True, MagicMock())]},
                ),
            ):
                with caplog.at_level("ERROR"):
                    generate_outputs_task(
                        scan_id=self.scan_id,
                        provider_id=self.provider_id,
                        tenant_id=self.tenant_id,
                    )
                    assert "Error deleting output files" in caplog.text

    @patch("tasks.tasks.rls_transaction")
    @patch("tasks.tasks.Integration.objects.filter")
    def test_generate_outputs_filters_enabled_s3_integrations(
        self, mock_integration_filter, mock_rls
    ):
        """Test that generate_outputs_task only processes enabled S3 integrations."""
        with (
            patch("tasks.tasks.ScanSummary.objects.filter") as mock_summary,
            patch("tasks.tasks.Provider.objects.get"),
            patch("tasks.tasks.initialize_prowler_provider"),
            patch("tasks.tasks.Compliance.get_bulk"),
            patch("tasks.tasks.get_compliance_frameworks", return_value=[]),
            patch("tasks.tasks.Finding.all_objects.filter") as mock_findings,
            patch(
                "tasks.tasks._generate_output_directory", return_value=("out", "comp")
            ),
            patch("tasks.tasks.FindingOutput._transform_findings_stats"),
            patch("tasks.tasks.FindingOutput.transform_api_finding"),
            patch("tasks.tasks._compress_output_files", return_value="/tmp/compressed"),
            patch("tasks.tasks._upload_to_s3", return_value="s3://bucket/file.zip"),
            patch("tasks.tasks.Scan.all_objects.filter"),
            patch("tasks.tasks.rmtree"),
            patch("tasks.tasks.s3_integration_task.apply_async") as mock_s3_task,
        ):
            mock_summary.return_value.exists.return_value = True
            mock_findings.return_value.order_by.return_value.iterator.return_value = [
                [MagicMock()],
                True,
            ]
            mock_integration_filter.return_value = [MagicMock()]
            mock_rls.return_value.__enter__.return_value = None

            with (
                patch("tasks.tasks.OUTPUT_FORMATS_MAPPING", {}),
                patch("tasks.tasks.COMPLIANCE_CLASS_MAP", {"aws": []}),
            ):
                generate_outputs_task(
                    scan_id=self.scan_id,
                    provider_id=self.provider_id,
                    tenant_id=self.tenant_id,
                )

            # Verify the S3 integrations filters
            mock_integration_filter.assert_called_once_with(
                integrationproviderrelationship__provider_id=self.provider_id,
                integration_type=Integration.IntegrationChoices.AMAZON_S3,
                enabled=True,
            )
            mock_s3_task.assert_called_once()


class TestScanCompleteTasks:
    @patch("tasks.tasks.create_compliance_requirements_task.apply_async")
    @patch("tasks.tasks.perform_scan_summary_task.si")
    @patch("tasks.tasks.generate_outputs_task.si")
    def test_scan_complete_tasks(
        self, mock_outputs_task, mock_scan_summary_task, mock_compliance_tasks
    ):
        _perform_scan_complete_tasks("tenant-id", "scan-id", "provider-id")
        mock_compliance_tasks.assert_called_once_with(
            kwargs={"tenant_id": "tenant-id", "scan_id": "scan-id"},
        )
        mock_scan_summary_task.assert_called_once_with(
            scan_id="scan-id",
            tenant_id="tenant-id",
        )
        mock_outputs_task.assert_called_once_with(
            scan_id="scan-id",
            provider_id="provider-id",
            tenant_id="tenant-id",
        )


@pytest.mark.django_db
class TestCheckIntegrationsTask:
    def setup_method(self):
        self.scan_id = str(uuid.uuid4())
        self.provider_id = str(uuid.uuid4())
        self.tenant_id = str(uuid.uuid4())
        self.output_directory = "/tmp/some-output-dir"

    @patch("tasks.tasks.rls_transaction")
    @patch("tasks.tasks.Integration.objects.filter")
    def test_check_integrations_no_integrations(
        self, mock_integration_filter, mock_rls
    ):
        mock_integration_filter.return_value.exists.return_value = False
        # Ensure rls_transaction is mocked
        mock_rls.return_value.__enter__.return_value = None

        result = check_integrations_task(
            tenant_id=self.tenant_id,
            provider_id=self.provider_id,
        )

        assert result == {"integrations_processed": 0}
        mock_integration_filter.assert_called_once_with(
            integrationproviderrelationship__provider_id=self.provider_id,
            enabled=True,
        )

    @patch("tasks.tasks.security_hub_integration_task")
    @patch("tasks.tasks.group")
    @patch("tasks.tasks.rls_transaction")
    @patch("tasks.tasks.Integration.objects.filter")
    def test_check_integrations_security_hub_success(
        self, mock_integration_filter, mock_rls, mock_group, mock_security_hub_task
    ):
        """Test that SecurityHub integrations are processed correctly."""
        # Mock that we have SecurityHub integrations
        mock_integrations = MagicMock()
        mock_integrations.exists.return_value = True

        # Mock SecurityHub integrations to return existing integrations
        mock_security_hub_integrations = MagicMock()
        mock_security_hub_integrations.exists.return_value = True

        # Set up the filter chain
        mock_integration_filter.return_value = mock_integrations
        mock_integrations.filter.return_value = mock_security_hub_integrations

        # Mock the task signature
        mock_task_signature = MagicMock()
        mock_security_hub_task.s.return_value = mock_task_signature

        # Mock group job
        mock_job = MagicMock()
        mock_group.return_value = mock_job

        # Ensure rls_transaction is mocked
        mock_rls.return_value.__enter__.return_value = None

        # Execute the function
        result = check_integrations_task(
            tenant_id=self.tenant_id,
            provider_id=self.provider_id,
            scan_id="test-scan-id",
        )

        # Should process 1 SecurityHub integration
        assert result == {"integrations_processed": 1}

        # Verify the integration filter was called
        mock_integration_filter.assert_called_once_with(
            integrationproviderrelationship__provider_id=self.provider_id,
            enabled=True,
        )

        # Verify SecurityHub integrations were filtered
        mock_integrations.filter.assert_called_once_with(
            integration_type=Integration.IntegrationChoices.AWS_SECURITY_HUB
        )

        # Verify SecurityHub task was created with correct parameters
        mock_security_hub_task.s.assert_called_once_with(
            tenant_id=self.tenant_id,
            provider_id=self.provider_id,
            scan_id="test-scan-id",
        )

        # Verify group was called and job was executed
        mock_group.assert_called_once_with([mock_task_signature])
        mock_job.apply_async.assert_called_once()

    @patch("tasks.tasks.rls_transaction")
    @patch("tasks.tasks.Integration.objects.filter")
    def test_check_integrations_disabled_integrations_ignored(
        self, mock_integration_filter, mock_rls
    ):
        """Test that disabled integrations are not processed."""
        mock_integration_filter.return_value.exists.return_value = False
        mock_rls.return_value.__enter__.return_value = None

        result = check_integrations_task(
            tenant_id=self.tenant_id,
            provider_id=self.provider_id,
        )

        assert result == {"integrations_processed": 0}
        mock_integration_filter.assert_called_once_with(
            integrationproviderrelationship__provider_id=self.provider_id,
            enabled=True,
        )

    @patch("tasks.tasks.s3_integration_task")
    @patch("tasks.tasks.Integration.objects.filter")
    @patch("tasks.tasks.ScanSummary.objects.filter")
    @patch("tasks.tasks.Provider.objects.get")
    @patch("tasks.tasks.initialize_prowler_provider")
    @patch("tasks.tasks.Compliance.get_bulk")
    @patch("tasks.tasks.get_compliance_frameworks")
    @patch("tasks.tasks.Finding.all_objects.filter")
    @patch("tasks.tasks._generate_output_directory")
    @patch("tasks.tasks.FindingOutput._transform_findings_stats")
    @patch("tasks.tasks.FindingOutput.transform_api_finding")
    @patch("tasks.tasks._compress_output_files")
    @patch("tasks.tasks._upload_to_s3")
    @patch("tasks.tasks.Scan.all_objects.filter")
    @patch("tasks.tasks.rmtree")
    def test_generate_outputs_with_asff_for_aws_with_security_hub(
        self,
        mock_rmtree,
        mock_scan_update,
        mock_upload,
        mock_compress,
        mock_transform_finding,
        mock_transform_stats,
        mock_generate_dir,
        mock_findings,
        mock_get_frameworks,
        mock_compliance_bulk,
        mock_initialize_provider,
        mock_provider_get,
        mock_scan_summary,
        mock_integration_filter,
        mock_s3_task,
    ):
        """Test that ASFF output is generated for AWS providers with SecurityHub integration."""
        # Setup
        mock_scan_summary_qs = MagicMock()
        mock_scan_summary_qs.exists.return_value = True
        mock_scan_summary.return_value = mock_scan_summary_qs

        # Mock AWS provider
        mock_provider = MagicMock()
        mock_provider.uid = "aws-account-123"
        mock_provider.provider = "aws"
        mock_provider_get.return_value = mock_provider

        # Mock SecurityHub integration exists
        mock_security_hub_integrations = MagicMock()
        mock_security_hub_integrations.exists.return_value = True
        mock_integration_filter.return_value = mock_security_hub_integrations

        # Mock s3_integration_task
        mock_s3_task.apply_async.return_value.get.return_value = True

        # Mock other necessary components
        mock_initialize_provider.return_value = MagicMock()
        mock_compliance_bulk.return_value = {}
        mock_get_frameworks.return_value = []
        mock_generate_dir.return_value = ("out-dir", "comp-dir")
        mock_transform_stats.return_value = {"stats": "data"}

        # Mock findings
        mock_finding = MagicMock()
        mock_findings.return_value.order_by.return_value.iterator.return_value = [
            [mock_finding],
            True,
        ]
        mock_transform_finding.return_value = MagicMock(compliance={})

        # Track which output formats were created
        created_writers = {}

        def track_writer_creation(cls_type):
            def factory(*args, **kwargs):
                writer = MagicMock()
                writer._data = []
                writer.transform = MagicMock()
                writer.batch_write_data_to_file = MagicMock()
                created_writers[cls_type] = writer
                return writer

            return factory

        # Mock OUTPUT_FORMATS_MAPPING with tracking
        with patch(
            "tasks.tasks.OUTPUT_FORMATS_MAPPING",
            {
                "csv": {
                    "class": track_writer_creation("csv"),
                    "suffix": ".csv",
                    "kwargs": {},
                },
                "json-asff": {
                    "class": track_writer_creation("asff"),
                    "suffix": ".asff.json",
                    "kwargs": {},
                },
                "json-ocsf": {
                    "class": track_writer_creation("ocsf"),
                    "suffix": ".ocsf.json",
                    "kwargs": {},
                },
            },
        ):
            mock_compress.return_value = "/tmp/compressed.zip"
            mock_upload.return_value = "s3://bucket/file.zip"

            # Execute
            result = generate_outputs_task(
                scan_id=self.scan_id,
                provider_id=self.provider_id,
                tenant_id=self.tenant_id,
            )

            # Verify ASFF was created for AWS with SecurityHub
            assert "asff" in created_writers, "ASFF writer should be created"
            assert "csv" in created_writers, "CSV writer should be created"
            assert "ocsf" in created_writers, "OCSF writer should be created"

            # Verify SecurityHub integration was checked
            assert mock_integration_filter.call_count == 2
            mock_integration_filter.assert_any_call(
                integrationproviderrelationship__provider_id=self.provider_id,
                integration_type=Integration.IntegrationChoices.AWS_SECURITY_HUB,
                enabled=True,
            )

            assert result == {"upload": True}

    @patch("tasks.tasks.s3_integration_task")
    @patch("tasks.tasks.Integration.objects.filter")
    @patch("tasks.tasks.ScanSummary.objects.filter")
    @patch("tasks.tasks.Provider.objects.get")
    @patch("tasks.tasks.initialize_prowler_provider")
    @patch("tasks.tasks.Compliance.get_bulk")
    @patch("tasks.tasks.get_compliance_frameworks")
    @patch("tasks.tasks.Finding.all_objects.filter")
    @patch("tasks.tasks._generate_output_directory")
    @patch("tasks.tasks.FindingOutput._transform_findings_stats")
    @patch("tasks.tasks.FindingOutput.transform_api_finding")
    @patch("tasks.tasks._compress_output_files")
    @patch("tasks.tasks._upload_to_s3")
    @patch("tasks.tasks.Scan.all_objects.filter")
    @patch("tasks.tasks.rmtree")
    def test_generate_outputs_no_asff_for_aws_without_security_hub(
        self,
        mock_rmtree,
        mock_scan_update,
        mock_upload,
        mock_compress,
        mock_transform_finding,
        mock_transform_stats,
        mock_generate_dir,
        mock_findings,
        mock_get_frameworks,
        mock_compliance_bulk,
        mock_initialize_provider,
        mock_provider_get,
        mock_scan_summary,
        mock_integration_filter,
        mock_s3_task,
    ):
        """Test that ASFF output is NOT generated for AWS providers without SecurityHub integration."""
        # Setup
        mock_scan_summary_qs = MagicMock()
        mock_scan_summary_qs.exists.return_value = True
        mock_scan_summary.return_value = mock_scan_summary_qs

        # Mock AWS provider
        mock_provider = MagicMock()
        mock_provider.uid = "aws-account-123"
        mock_provider.provider = "aws"
        mock_provider_get.return_value = mock_provider

        # Mock NO SecurityHub integration
        mock_security_hub_integrations = MagicMock()
        mock_security_hub_integrations.exists.return_value = False
        mock_integration_filter.return_value = mock_security_hub_integrations

        # Mock other necessary components
        mock_initialize_provider.return_value = MagicMock()
        mock_compliance_bulk.return_value = {}
        mock_get_frameworks.return_value = []
        mock_generate_dir.return_value = ("out-dir", "comp-dir")
        mock_transform_stats.return_value = {"stats": "data"}

        # Mock findings
        mock_finding = MagicMock()
        mock_findings.return_value.order_by.return_value.iterator.return_value = [
            [mock_finding],
            True,
        ]
        mock_transform_finding.return_value = MagicMock(compliance={})

        # Track which output formats were created
        created_writers = {}

        def track_writer_creation(cls_type):
            def factory(*args, **kwargs):
                writer = MagicMock()
                writer._data = []
                writer.transform = MagicMock()
                writer.batch_write_data_to_file = MagicMock()
                created_writers[cls_type] = writer
                return writer

            return factory

        # Mock OUTPUT_FORMATS_MAPPING with tracking
        with patch(
            "tasks.tasks.OUTPUT_FORMATS_MAPPING",
            {
                "csv": {
                    "class": track_writer_creation("csv"),
                    "suffix": ".csv",
                    "kwargs": {},
                },
                "json-asff": {
                    "class": track_writer_creation("asff"),
                    "suffix": ".asff.json",
                    "kwargs": {},
                },
                "json-ocsf": {
                    "class": track_writer_creation("ocsf"),
                    "suffix": ".ocsf.json",
                    "kwargs": {},
                },
            },
        ):
            mock_compress.return_value = "/tmp/compressed.zip"
            mock_upload.return_value = "s3://bucket/file.zip"

            # Execute
            result = generate_outputs_task(
                scan_id=self.scan_id,
                provider_id=self.provider_id,
                tenant_id=self.tenant_id,
            )

            # Verify ASFF was NOT created when no SecurityHub integration
            assert "asff" not in created_writers, "ASFF writer should NOT be created"
            assert "csv" in created_writers, "CSV writer should be created"
            assert "ocsf" in created_writers, "OCSF writer should be created"

            # Verify SecurityHub integration was checked
            assert mock_integration_filter.call_count == 2
            mock_integration_filter.assert_any_call(
                integrationproviderrelationship__provider_id=self.provider_id,
                integration_type=Integration.IntegrationChoices.AWS_SECURITY_HUB,
                enabled=True,
            )

            assert result == {"upload": True}

    @patch("tasks.tasks.ScanSummary.objects.filter")
    @patch("tasks.tasks.Provider.objects.get")
    @patch("tasks.tasks.initialize_prowler_provider")
    @patch("tasks.tasks.Compliance.get_bulk")
    @patch("tasks.tasks.get_compliance_frameworks")
    @patch("tasks.tasks.Finding.all_objects.filter")
    @patch("tasks.tasks._generate_output_directory")
    @patch("tasks.tasks.FindingOutput._transform_findings_stats")
    @patch("tasks.tasks.FindingOutput.transform_api_finding")
    @patch("tasks.tasks._compress_output_files")
    @patch("tasks.tasks._upload_to_s3")
    @patch("tasks.tasks.Scan.all_objects.filter")
    @patch("tasks.tasks.rmtree")
    def test_generate_outputs_no_asff_for_non_aws_provider(
        self,
        mock_rmtree,
        mock_scan_update,
        mock_upload,
        mock_compress,
        mock_transform_finding,
        mock_transform_stats,
        mock_generate_dir,
        mock_findings,
        mock_get_frameworks,
        mock_compliance_bulk,
        mock_initialize_provider,
        mock_provider_get,
        mock_scan_summary,
    ):
        """Test that ASFF output is NOT generated for non-AWS providers (e.g., Azure, GCP)."""
        # Setup
        mock_scan_summary_qs = MagicMock()
        mock_scan_summary_qs.exists.return_value = True
        mock_scan_summary.return_value = mock_scan_summary_qs

        # Mock Azure provider (non-AWS)
        mock_provider = MagicMock()
        mock_provider.uid = "azure-subscription-123"
        mock_provider.provider = "azure"  # Non-AWS provider
        mock_provider_get.return_value = mock_provider

        # Mock other necessary components
        mock_initialize_provider.return_value = MagicMock()
        mock_compliance_bulk.return_value = {}
        mock_get_frameworks.return_value = []
        mock_generate_dir.return_value = ("out-dir", "comp-dir")
        mock_transform_stats.return_value = {"stats": "data"}

        # Mock findings
        mock_finding = MagicMock()
        mock_findings.return_value.order_by.return_value.iterator.return_value = [
            [mock_finding],
            True,
        ]
        mock_transform_finding.return_value = MagicMock(compliance={})

        # Track which output formats were created
        created_writers = {}

        def track_writer_creation(cls_type):
            def factory(*args, **kwargs):
                writer = MagicMock()
                writer._data = []
                writer.transform = MagicMock()
                writer.batch_write_data_to_file = MagicMock()
                created_writers[cls_type] = writer
                return writer

            return factory

        # Mock OUTPUT_FORMATS_MAPPING with tracking
        with patch(
            "tasks.tasks.OUTPUT_FORMATS_MAPPING",
            {
                "csv": {
                    "class": track_writer_creation("csv"),
                    "suffix": ".csv",
                    "kwargs": {},
                },
                "json-asff": {
                    "class": track_writer_creation("asff"),
                    "suffix": ".asff.json",
                    "kwargs": {},
                },
                "json-ocsf": {
                    "class": track_writer_creation("ocsf"),
                    "suffix": ".ocsf.json",
                    "kwargs": {},
                },
            },
        ):
            mock_compress.return_value = "/tmp/compressed.zip"
            mock_upload.return_value = "s3://bucket/file.zip"

            # Execute
            result = generate_outputs_task(
                scan_id=self.scan_id,
                provider_id=self.provider_id,
                tenant_id=self.tenant_id,
            )

            # Verify ASFF was NOT created for non-AWS provider
            assert (
                "asff" not in created_writers
            ), "ASFF writer should NOT be created for non-AWS providers"
            assert "csv" in created_writers, "CSV writer should be created"
            assert "ocsf" in created_writers, "OCSF writer should be created"

            assert result == {"upload": True}

    @patch("tasks.tasks.upload_s3_integration")
    def test_s3_integration_task_success(self, mock_upload):
        mock_upload.return_value = True
        output_directory = "/tmp/prowler_api_output/test"

        result = s3_integration_task(
            tenant_id=self.tenant_id,
            provider_id=self.provider_id,
            output_directory=output_directory,
        )

        assert result is True
        mock_upload.assert_called_once_with(
            self.tenant_id, self.provider_id, output_directory
        )

    @patch("tasks.tasks.upload_s3_integration")
    def test_s3_integration_task_failure(self, mock_upload):
        mock_upload.return_value = False
        output_directory = "/tmp/prowler_api_output/test"

        result = s3_integration_task(
            tenant_id=self.tenant_id,
            provider_id=self.provider_id,
            output_directory=output_directory,
        )

        assert result is False
        mock_upload.assert_called_once_with(
            self.tenant_id, self.provider_id, output_directory
        )

    @patch("tasks.tasks.upload_security_hub_integration")
    def test_security_hub_integration_task_success(self, mock_upload):
        """Test successful SecurityHub integration task execution."""
        mock_upload.return_value = True
        scan_id = "test-scan-123"

        result = security_hub_integration_task(
            tenant_id=self.tenant_id,
            provider_id=self.provider_id,
            scan_id=scan_id,
        )

        assert result is True
        mock_upload.assert_called_once_with(self.tenant_id, self.provider_id, scan_id)

    @patch("tasks.tasks.upload_security_hub_integration")
    def test_security_hub_integration_task_failure(self, mock_upload):
        """Test SecurityHub integration task handling failure."""
        mock_upload.return_value = False
        scan_id = "test-scan-123"

        result = security_hub_integration_task(
            tenant_id=self.tenant_id,
            provider_id=self.provider_id,
            scan_id=scan_id,
        )

        assert result is False
        mock_upload.assert_called_once_with(self.tenant_id, self.provider_id, scan_id)
