test(exports): remove time.sleep() from unit concurrency tests

Author: noliveleger

kobo/settings/base.py

@@ -477,7 +477,7 @@
     ),
     'ANONYMOUS_EXPORTS_GRACE_PERIOD': (
         30,
-        'Number of minutes after which anonymous export tasks are cleaned up.'
+        'Number of minutes after which anonymous export tasks are cleaned up.',
     ),
     'LIMIT_ATTACHMENT_REMOVAL_GRACE_PERIOD': (
         90,
@@ -1453,7 +1453,7 @@ def dj_stripe_request_callback_method():
     'cleanup-anonymous-exports': {
         'task': 'kpi.tasks.cleanup_anonymous_exports',
         'schedule': crontab(minute='*/5'),
-        'options': {'queue': 'kpi_low_priority_queue'}
+        'options': {'queue': 'kpi_low_priority_queue'},
     },
     # Schedule every 15 minutes
     'refresh-user-report-snapshot': {

kpi/tasks.py

@@ -89,12 +89,15 @@ def cleanup_anonymous_exports(**kwargs):
         minutes=config.ANONYMOUS_EXPORTS_GRACE_PERIOD
     )
 
-    old_export_ids = SubmissionExportTask.objects.filter(
-        user=get_anonymous_user(),
-        date_created__lt=cutoff_time,
-    ).exclude(
-        status=ImportExportStatusChoices.PROCESSING
-    ).order_by('date_created').values_list('pk', flat=True)[:BATCH_SIZE]
+    old_export_ids = (
+        SubmissionExportTask.objects.filter(
+            user=get_anonymous_user(),
+            date_created__lt=cutoff_time,
+        )
+        .exclude(status=ImportExportStatusChoices.PROCESSING)
+        .order_by('date_created')
+        .values_list('pk', flat=True)[:BATCH_SIZE]
+    )
 
     if not old_export_ids:
         logging.info('No old anonymous exports to clean up.')
@@ -104,7 +107,7 @@ def cleanup_anonymous_exports(**kwargs):
     for export_id in old_export_ids:
         try:
             with transaction.atomic():
-                # Acquire a row level lock without waiting
+                # Acquire a row-level lock without waiting
                 export = (
                     SubmissionExportTask.objects.only('pk', 'uid', 'result')
                     .select_for_update(nowait=True)
@@ -121,9 +124,7 @@ def cleanup_anonymous_exports(**kwargs):
                 export.delete()
                 deleted_count += 1
         except DatabaseError:
-            logging.info(
-                f'Export {export_id} is currently being processed. Skipping.'
-            )
+            logging.info(f'Export {export_id} is currently being processed. Skipping.')
     logging.info(f'Cleaned up {deleted_count} old anonymous exports.')
 
 

kpi/tests/test_cleanup_anonymous_exports.py

@@ -1,14 +1,13 @@
 import threading
-import time
 from datetime import timedelta
 
 from django.db import transaction
-from django.utils import timezone
 from django.test import TransactionTestCase
+from django.utils import timezone
 
 from kpi.models.import_export_task import (
     ImportExportStatusChoices,
-    SubmissionExportTask
+    SubmissionExportTask,
 )
 from kpi.tasks import cleanup_anonymous_exports
 from kpi.utils.object_permission import get_anonymous_user
@@ -52,62 +51,52 @@ def test_processing_exports_are_not_deleted(self):
         Test that exports with PROCESSING status are never deleted
         """
         processing_export = self._create_export_task(
-            status=ImportExportStatusChoices.PROCESSING,
-            minutes_old=100
+            status=ImportExportStatusChoices.PROCESSING, minutes_old=100
         )
 
         cleanup_anonymous_exports()
         self.assertTrue(
-            SubmissionExportTask.objects.filter(
-                uid=processing_export.uid
-            ).exists()
+            SubmissionExportTask.objects.filter(uid=processing_export.uid).exists()
         )
 
     def test_concurrency_locked_rows_are_skipped(self):
         exports = [self._create_export_task(minutes_old=120) for _ in range(5)]
-        export_pks = [e.pk for e in exports]
-        locked_pks = export_pks[:3]
-        unlocked_pks = export_pks[3:]
-
-        def lock_and_hold():
-            """
-            Acquire lock on first 3 exports and hold for 5 seconds,
-            this will block cleanup from acquiring the lock on these rows and
-            in the meantime other rows should be cleaned up
-            """
+        locked_export = exports[0]
+        not_locked_exports = exports[1:]
+
+        lock_acquired = threading.Event()
+        release_lock = threading.Event()
+
+        def lock_row():
             with transaction.atomic():
-                # Acquire lock on first 3 exports
-                list(
-                    SubmissionExportTask.objects
-                    .select_for_update()
-                    .filter(pk__in=locked_pks)
+                # Lock the row
+                SubmissionExportTask.objects.select_for_update().get(
+                    pk=locked_export.pk
                 )
 
-                # Hold lock for 5 seconds,
-                # during this time cleanup should run
-                time.sleep(5)
+                # Signal that the lock is active
+                lock_acquired.set()
+
+                # Wait until the main thread signals we can release the lock
+                release_lock.wait()
 
-        # Start thread that will hold the lock
-        lock_thread = threading.Thread(target=lock_and_hold, daemon=True)
-        lock_thread.start()
+        t = threading.Thread(target=lock_row)
+        t.start()
 
-        # Give thread time to acquire lock
-        time.sleep(1)
+        # Wait for the row-level lock to actually be acquired
+        lock_acquired.wait()
 
-        # Run cleanup while lock is held
+        # Now cleanup should try select_for_update(nowait=True), causing DatabaseError
         cleanup_anonymous_exports()
 
-        # Wait for thread to finish
-        lock_thread.join(timeout=10)
+        # Let the locking thread finish its transaction
+        release_lock.set()
+        t.join()
 
-        # Verify locked rows were not deleted
-        remaining_locked = SubmissionExportTask.objects.filter(
-            pk__in=locked_pks
-        ).count()
-        self.assertEqual(remaining_locked, 3)
+        # Verify the locked row was not deleted
+        assert SubmissionExportTask.objects.filter(pk=locked_export.pk).exists()
 
         # Verify unlocked rows were deleted
-        remaining_unlocked = SubmissionExportTask.objects.filter(
-            pk__in=unlocked_pks
-        ).count()
-        self.assertEqual(remaining_unlocked, 0)
+        assert not SubmissionExportTask.objects.filter(
+            pk__in=[e.pk for e in not_locked_exports]
+        ).exists()