Fix tests

Signed-off-by: Tushar Goel <[email protected]>

Author: TG1999

vulnerabilities/models.py

@@ -1360,7 +1360,8 @@ class Meta:
 
     def save(self, *args, **kwargs):
         advisory_data = self.to_advisory_data()
-        self.unique_content_id = compute_content_id(advisory_data, include_metadata=False)
+        if not self.unique_content_id:
+            self.unique_content_id = compute_content_id(advisory_data, include_metadata=False)
         super().save(*args, **kwargs)
 
     def to_advisory_data(self) -> "AdvisoryData":

vulnerabilities/pipelines/recompute_content_ids.py

@@ -92,14 +92,14 @@ def process_advisories(
         for advisory_ids in progress.iter(advisory_batches):
             progress.log_progress()
             logger.debug(f"{advisory_func.__name__} len={len(advisory_ids)}")
-            advisory_func(advisory_ids=advisory_ids, logger=None)
+            advisory_func(advisory_ids=advisory_ids, logger=progress_logger)
         return
 
     logger.info(f"Starting ProcessPoolExecutor with {max_workers} max_workers")
 
     with futures.ProcessPoolExecutor(max_workers) as executor:
         future_to_advisories = {
-            executor.submit(advisory_func, advisory_ids, None): advisory_ids
+            executor.submit(advisory_func, advisory_ids, progress_logger): advisory_ids
             for advisory_ids in advisory_batches
         }
 
@@ -187,6 +187,7 @@ def recompute_content_ids(self):
         """
         while True:
             advisories = Advisory.objects.exclude(unique_content_id__length=64)
+            print(f"advisories: {advisories.count()}")
             if not advisories.exists():
                 break
             process_advisories(

vulnerabilities/pipelines/remove_duplicate_advisories.py

@@ -27,7 +27,7 @@ def remove_duplicates_batch(advisory_ids, logger=None):
     try:
         with transaction.atomic():
             advisories = Advisory.objects.filter(id__in=advisory_ids).select_for_update(
-                nowait=True, skip_locked=True
+                skip_locked=True
             )
             if not advisories.exists():
                 return
@@ -45,28 +45,25 @@ def remove_duplicates_batch(advisory_ids, logger=None):
                 if len(group_advisories) <= 1:
                     continue
 
-                if logger:
-                    logger.info(
-                        f"Found {len(group_advisories)} duplicates for content ID {content_id}",
-                    )
+                logger(
+                    f"Found {len(group_advisories)} duplicates for content ID {content_id}",
+                )
 
                 oldest = min(group_advisories, key=lambda x: x.date_imported)
 
                 advisory_ids_to_delete = [adv.id for adv in group_advisories if adv.id != oldest.id]
                 if advisory_ids_to_delete:
                     Advisory.objects.filter(id__in=advisory_ids_to_delete).delete()
-                    if logger:
-                        logger.info(
-                            f"Kept advisory {oldest.id} and removed "
-                            f"{len(advisory_ids_to_delete)} duplicates for content ID {content_id}",
-                        )
+                    logger(
+                        f"Kept advisory {oldest.id} and removed "
+                        f"{len(advisory_ids_to_delete)} duplicates for content ID {content_id}",
+                    )
 
     except Exception as e:
-        if logger:
-            logger(
-                f"Error processing batch of advisories: {e}",
-                level=logging.ERROR,
-            )
+        logger(
+            f"Error processing batch of advisories: {e}",
+            level=logging.ERROR,
+        )
 
 
 class RemoveDuplicateAdvisoriesPipeline(VulnerableCodePipeline):
@@ -91,19 +88,14 @@ def remove_duplicates(self):
                 .filter(count__gt=1)
                 .values_list("unique_content_id", flat=True)
             )
-
-            print(f"duplicate_content_ids: {duplicate_content_ids}")
-
             advisories = Advisory.objects.filter(unique_content_id__in=duplicate_content_ids)
-
             if not advisories.exists():
                 break
 
             self.log(
                 f"Processing {advisories.count()} content IDs with duplicates",
                 level=logging.INFO,
             )
-
             process_advisories(
                 advisories=advisories,
                 advisory_func=remove_duplicates_batch,

vulnerabilities/tests/test_recompute_content_id.py

@@ -44,7 +44,7 @@ def advisory_data():
 
 
 @pytest.mark.django_db
-def test_recompute_content_ids_basic(advisory_data):
+def test_recompute_content_ids_basic_async(advisory_data):
     """
     Test that advisories without content IDs get them computed.
     """
@@ -68,7 +68,7 @@ def test_recompute_content_ids_basic(advisory_data):
 
 
 @pytest.mark.django_db
-def test_recompute_content_ids_multiple_batches(advisory_data):
+def test_recompute_content_ids_multiple_batches_async(advisory_data):
     """
     Test that content ID computation works across multiple batches.
     """
@@ -103,6 +103,66 @@ def test_recompute_content_ids_multiple_batches(advisory_data):
     assert Advisory.objects.exclude(unique_content_id__length=64).count() == 0
 
 
+@pytest.mark.django_db
+def test_recompute_content_ids_basic(advisory_data):
+    """
+    Test that advisories without content IDs get them computed.
+    """
+    advisory = Advisory.objects.create(
+        summary=advisory_data.summary,
+        affected_packages=[pkg.to_dict() for pkg in advisory_data.affected_packages],
+        references=[ref.to_dict() for ref in advisory_data.references],
+        unique_content_id="",
+        date_collected=datetime.datetime(2024, 1, 1, tzinfo=pytz.UTC),
+    )
+
+    with patch("vulnerabilities.pipelines.recompute_content_ids.get_max_workers") as mock_workers:
+        mock_workers.return_value = 0
+
+        pipeline = RecomputeContentIDPipeline()
+        pipeline.recompute_content_ids()
+
+    advisory.refresh_from_db()
+    assert advisory.unique_content_id != ""
+    assert len(advisory.unique_content_id) == 64  # SHA256 hash length
+
+
+@pytest.mark.django_db
+def test_recompute_content_ids_multiple_batches(advisory_data):
+    """
+    Test that content ID computation works across multiple batches.
+    """
+    dates = [
+        datetime.datetime(
+            2024 + (i // (12 * 28)),  # Year
+            ((i // 28) % 12) + 1,  # Month (1-12)
+            (i % 28) + 1,  # Day (1-28)
+            tzinfo=pytz.UTC,
+        )
+        for i in range(2500)  # Create 2500 advisories
+    ]
+
+    for date in dates:
+        Advisory.objects.create(
+            summary=advisory_data.summary,
+            affected_packages=[pkg.to_dict() for pkg in advisory_data.affected_packages],
+            references=[ref.to_dict() for ref in advisory_data.references],
+            unique_content_id="",
+            date_imported=date,
+            date_collected=date,
+        )
+
+    with patch("vulnerabilities.pipelines.recompute_content_ids.get_max_workers") as mock_workers:
+        mock_workers.return_value = 0
+
+        pipeline = RecomputeContentIDPipeline()
+        pipeline.BATCH_SIZE = 1000
+        pipeline.recompute_content_ids()
+
+    assert not Advisory.objects.filter(unique_content_id="").exists()
+    assert Advisory.objects.exclude(unique_content_id__length=64).count() == 0
+
+
 @pytest.mark.django_db
 def test_recompute_content_ids_preserves_existing(advisory_data):
     """
@@ -137,7 +197,7 @@ def test_recompute_content_ids_error_handling(advisory_data):
         summary=advisory_data.summary,
         affected_packages=[pkg.to_dict() for pkg in advisory_data.affected_packages],
         references=[ref.to_dict() for ref in advisory_data.references],
-        unique_content_id="",
+        unique_content_id="Test",
         date_collected=datetime.datetime(2024, 1, 1, tzinfo=pytz.UTC),
     )
 
@@ -150,8 +210,6 @@ def test_recompute_content_ids_error_handling(advisory_data):
             mock_workers.return_value = 0
 
             pipeline = RecomputeContentIDPipeline()
-            pipeline.recompute_content_ids()
-
-    advisory = Advisory.objects.first()
-    assert advisory is not None
-    assert advisory.unique_content_id == ""
+            # expect an error
+            with pytest.raises(Exception):
+                pipeline.recompute_content_ids()

vulnerabilities/tests/test_remove_duplicate_advisories.py

@@ -65,8 +65,39 @@ def test_remove_duplicates_keeps_oldest(advisory_data):
         )
 
     with patch("vulnerabilities.pipelines.recompute_content_ids.get_max_workers") as mock_workers:
-        mock_workers.return_value = 4  # Simulate 4 workers with keep_available=0
+        mock_workers.return_value = 0  # Simulate 4 workers with keep_available=0
+        pipeline = RemoveDuplicateAdvisoriesPipeline()
+        pipeline.remove_duplicates()
+
+    # Check that only the oldest advisory remains
+    remaining = Advisory.objects.all()
+    assert remaining.count() == 1
+    assert remaining.first().date_imported == dates[0]
 
+
+@pytest.mark.django_db
+def test_remove_duplicates_keeps_oldest_async(advisory_data):
+    """
+    Test that when multiple advisories have the same content,
+    only the oldest one is kept.
+    """
+    dates = [
+        datetime.datetime(2024, 1, 1, tzinfo=pytz.UTC),
+        datetime.datetime(2024, 1, 2, tzinfo=pytz.UTC),
+        datetime.datetime(2024, 1, 3, tzinfo=pytz.UTC),
+    ]
+
+    for date in dates:
+        Advisory.objects.create(
+            summary=advisory_data.summary,
+            affected_packages=[pkg.to_dict() for pkg in advisory_data.affected_packages],
+            references=[ref.to_dict() for ref in advisory_data.references],
+            date_imported=date,
+            date_collected=date,
+        )
+
+    with patch("vulnerabilities.pipelines.recompute_content_ids.get_max_workers") as mock_workers:
+        mock_workers.return_value = 4  # Simulate 4 workers with keep_available=0
         pipeline = RemoveDuplicateAdvisoriesPipeline()
         pipeline.remove_duplicates()
 
@@ -122,7 +153,7 @@ def test_remove_duplicates_with_multiple_batches(advisory_data):
             (i % 28) + 1,  # Day (1-28)
             tzinfo=pytz.UTC,
         )
-        for i in range(100)  # Create 2500 advisories
+        for i in range(100)
     ]
 
     for date in dates: