Fix tests

Signed-off-by: Tushar Goel <tushar.goel.dav@gmail.com>

Author: TG1999

…_advisory_latest_by_avid_idx_and_more.py …mpacts_unfurled_successfully_and_more.pyvulnerabilities/migrations/0135_remove_advisoryv2_advisory_latest_by_avid_idx_and_more.py renamed to vulnerabilities/migrations/0135_advisoryv2__all_impacts_unfurled_successfully_and_more.py vulnerabilities/migrations/0135_remove_advisoryv2_advisory_latest_by_avid_idx_and_more.py renamed to vulnerabilities/migrations/0135_advisoryv2__all_impacts_unfurled_successfully_and_more.py

@@ -1,4 +1,4 @@
-# Generated by Django 5.2.11 on 2026-06-03 08:00
+# Generated by Django 5.2.11 on 2026-06-03 12:46
 
 from django.db import migrations, models
 
@@ -10,9 +10,13 @@ class Migration(migrations.Migration):
     ]
 
     operations = [
-        migrations.RemoveIndex(
+        migrations.AddField(
             model_name="advisoryv2",
-            name="advisory_latest_by_avid_idx",
+            name="_all_impacts_unfurled_successfully",
+            field=models.BooleanField(
+                default=False,
+                help_text="Indicates whether all impacts for this advisory have been unfurled successfully.",
+            ),
         ),
         migrations.AddIndex(
             model_name="advisoryv2",

vulnerabilities/models.py

@@ -3258,6 +3258,11 @@ class AdvisoryV2(models.Model):
         help_text="Indicates whether all impacts for this advisory have been unfurled.",
     )
 
+    _all_impacts_unfurled_successfully = models.BooleanField(
+        default=False,
+        help_text="Indicates whether all impacts for this advisory have been unfurled successfully.",
+    )
+
     objects = AdvisoryV2QuerySet.as_manager()
 
     class Meta:
@@ -3272,9 +3277,7 @@ class Meta:
             models.Index(
                 fields=["avid", "-date_collected", "-id"],
                 name="advisory_latest_by_avid_idx",
-            )
-        ]
-        indexes = [
+            ),
             models.Index(
                 fields=["_all_impacts_unfurled", "id"],
                 name="advisory_unfurled_idx",

vulnerabilities/pipelines/v2_improvers/unfurl_version_range.py

@@ -51,10 +51,15 @@ class UnfurlVersionRangePipeline(VulnerableCodePipeline):
 
     # Days elapsed before version range is re-unfurled
     reunfurl_after_days = 2
+    impacted_packages = None
 
     @classmethod
     def steps(cls):
-        return (cls.unfurl_version_range, cls.mark_all_impacts_unfurled)
+        return (
+            cls.unfurl_version_range,
+            cls.mark_all_impacts_unfurled_sucessfully,
+            cls.mark_all_impacts_unfurl_attempted,
+        )
 
     def unfurl_version_range(self):
         processed_impacted_packages_count = 0
@@ -66,6 +71,7 @@ def unfurl_version_range(self):
         chunk_size = 500
 
         impacted_packages = impacted_package_qs(cutoff_day=self.reunfurl_after_days)
+        self.impacted_packages = impacted_packages
         impacted_packages_count = impacted_packages.count()
         self.log(f"Unfurl affected vers range for {impacted_packages_count:,d} ImpactedPackage.")
 
@@ -104,32 +110,67 @@ def unfurl_version_range(self):
             processed_impacted_packages_count += 1
 
             if len(update_unfurl_date) > update_batch_size:
+                cur_time = timezone.now()
                 ImpactedPackage.objects.filter(pk__in=update_unfurl_date).update(
-                    last_range_unfurl_at=timezone.now()
+                    last_range_unfurl_at=cur_time
                 )
                 ImpactedPackage.objects.filter(pk__in=update_successful_unfurl_date).update(
-                    last_successful_range_unfurl_at=timezone.now()
+                    last_successful_range_unfurl_at=cur_time
                 )
                 update_unfurl_date.clear()
                 update_successful_unfurl_date.clear()
 
+        cur_time = timezone.now()
+
         ImpactedPackage.objects.filter(pk__in=update_unfurl_date).update(
-            last_range_unfurl_at=timezone.now()
+            last_range_unfurl_at=cur_time
         )
         ImpactedPackage.objects.filter(pk__in=update_successful_unfurl_date).update(
-            last_successful_range_unfurl_at=timezone.now()
+            last_successful_range_unfurl_at=cur_time
         )
         self.log(f"Successfully processed {processed_impacted_packages_count:,d} ImpactedPackage.")
         self.log(f"{processed_affected_packages_count:,d} new Impact-Package relation created.")
 
-    def mark_all_impacts_unfurled(self):
+    def mark_all_impacts_unfurled_sucessfully(self):
+        impacted_packages = self.impacted_packages or impacted_package_qs(
+            cutoff_day=self.reunfurl_after_days
+        )
         while True:
-            advisories = list(latest_advisories_with_all_impacts_unfurled()[:100])
+            advisory_ids = list(
+                latest_advisories_with_all_impacts_unfurled_successfully(
+                    impacted_packages=impacted_packages
+                )[:100]
+            )
 
-            if not advisories:
+            if not advisory_ids:
                 break
 
-            complete_advisories_import(AdvisoryV2.objects.filter(id__in=[a.id for a in advisories]))
+            complete_advisories_import(advisory_ids=advisory_ids, success=True)
+
+    def mark_all_impacts_unfurl_attempted(self):
+        impacted_packages = self.impacted_packages or impacted_package_qs(
+            cutoff_day=self.reunfurl_after_days
+        )
+        advisories_qs = latest_advisories_with_all_impacts_unfurled_attempted(
+            impacted_packages=impacted_packages
+        )
+
+        batch_size = 100
+        batch = []
+
+        for advisory_id in advisories_qs.iterator(chunk_size=100):
+            batch.append(advisory_id)
+
+            if len(batch) >= batch_size:
+                complete_advisories_import(
+                    advisory_ids=list(batch),
+                )
+                batch.clear()
+
+        if batch:
+            complete_advisories_import(
+                advisory_ids=list(batch),
+            )
 
 
 def get_affected_purls(versions, impact, logger):
@@ -219,15 +260,17 @@ def impacted_package_qs(cutoff_day=2):
 
 
 @transaction.atomic
-def complete_advisories_import(advisories):
-
-    advisory_ids = list(advisories.values_list("id", flat=True))
-
+def complete_advisories_import(advisory_ids, success=False):
     if not advisory_ids:
         return
 
     AdvisoryV2.objects.filter(id__in=advisory_ids).update(_all_impacts_unfurled=True)
 
+    if success:
+        AdvisoryV2.objects.filter(id__in=advisory_ids).update(
+            _all_impacts_unfurled_successfully=True
+        )
+
     affecting_package_ids = set(
         ImpactedPackageAffecting.objects.filter(
             impacted_package__advisory_id__in=advisory_ids
@@ -256,16 +299,42 @@ def complete_advisories_import(advisories):
         group_advisory_for_package(package)
 
 
-def latest_advisories_with_all_impacts_unfurled():
-    remaining_unfurled_impacts = ImpactedPackage.objects.filter(
+def latest_advisories_with_all_impacts_unfurled_successfully(
+    impacted_packages=None,
+):
+    unsuccessful_impacts = impacted_packages.filter(
+        advisory_id=OuterRef("pk"),
+        advisory__is_latest=True,
+    ).filter(Q(last_range_unfurl_at__isnull=True) | Q(last_successful_range_unfurl_at__isnull=True))
+
+    return (
+        AdvisoryV2.objects.filter(
+            _all_impacts_unfurled_successfully=False,
+            is_latest=True,
+        )
+        .annotate(has_unsuccessful_impacts=Exists(unsuccessful_impacts))
+        .filter(has_unsuccessful_impacts=False)
+        .order_by("id")
+        .values_list("id", flat=True)
+    )
+
+
+def latest_advisories_with_all_impacts_unfurled_attempted(
+    impacted_packages=None,
+):
+    impacts_not_attempted = impacted_packages.filter(
         advisory_id=OuterRef("pk"),
-        last_range_unfurl_at__isnull=True,
         advisory__is_latest=True,
+        last_range_unfurl_at__isnull=True,
     )
 
     return (
-        AdvisoryV2.objects.filter(_all_impacts_unfurled=False, is_latest=True)
-        .annotate(has_remaining_unfurled=Exists(remaining_unfurled_impacts))
-        .filter(has_remaining_unfurled=False)
+        AdvisoryV2.objects.filter(
+            _all_impacts_unfurled_successfully=False,
+            is_latest=True,
+        )
+        .annotate(has_unattempted_impacts=Exists(impacts_not_attempted))
+        .filter(has_unattempted_impacts=False)
         .order_by("id")
+        .values_list("id", flat=True)
     )

vulnerabilities/pipes/advisory.py

@@ -371,6 +371,7 @@ def insert_advisory_v2(
     advisory_obj.risk_score = round(risk_score, 1) if risk_score is not None else None
     if not advisory.affected_packages:
         advisory_obj._all_impacts_unfurled = True
+        advisory_obj._all_impacts_unfurled_successfully = True
     advisory_obj.save()
 
     for affected_pkg in advisory.affected_packages:

vulnerabilities/tests/pipelines/exporters/test_federate_vulnerabilities.py

@@ -83,18 +83,25 @@ def setUp(self):
             date_published=datetime.now() - timedelta(days=10),
             url="https://example.com/advisory/2",
         )
-        insert_advisory_v2(
+        a1 = insert_advisory_v2(
             advisory=advisory1,
             pipeline_id="test_pipeline_v2",
             logger=self.logger.write,
             datasource_id="test",
         )
-        insert_advisory_v2(
+        a1._all_impacts_unfurled = True
+        a1._all_impacts_unfurled_successfully = True
+        a1.save()
+        a2 = insert_advisory_v2(
             advisory=advisory2,
             pipeline_id="test_pipeline_v2",
             logger=self.logger.write,
             datasource_id="test",
         )
+        a2._all_impacts_unfurled = True
+        a2._all_impacts_unfurled_successfully = True
+        a2.save()
+
 
     @patch(
         "vulnerabilities.pipelines.exporters.federate_vulnerabilities.FederatePackageVulnerabilities.clone_federation_repository"

vulnerabilities/tests/pipelines/v2_improvers/test_compute_package_risk_v2.py

@@ -36,6 +36,7 @@ def test_simple_risk_pipeline():
         url="https://test.com",
         date_collected=datetime.now(),
         is_latest=True,
+        _all_impacts_unfurled=True
     )
     adv.save()