fix tests.share.test_celery

Author: aaxelb

share/celery.py

@@ -1,7 +1,7 @@
+import datetime
 import functools
 import logging
 
-
 from celery import states
 from celery.app.task import Context
 from celery.backends.base import BaseDictBackend
@@ -92,7 +92,10 @@ def _store_result(self, task_id, result, status, traceback=None, request=None, *
 
     @die_on_unhandled
     def cleanup(self, expires=None):
-        TaskResultCleaner(expires or self.expires).clean()
+        TaskResultCleaner(
+            success_ttl=(expires or self.expires),
+            nonsuccess_ttl=settings.FAILED_CELERY_RESULT_EXPIRES,
+        ).clean()
 
     @die_on_unhandled
     def _get_task_meta_for(self, task_id):
@@ -113,11 +116,19 @@ class TaskResultCleaner:
 
     TaskModel = CeleryTaskResult
 
-    def __init__(self, expires, bucket=None, delete=True, chunk_size=5000):
-        self.bucket = bucket
+    def __init__(self, success_ttl, nonsuccess_ttl=None, delete=True, chunk_size=5000):
         self.chunk_size = chunk_size
         self.delete = delete
-        self.expires = expires
+        self.success_ttl = success_ttl
+        self.nonsuccess_ttl = nonsuccess_ttl or success_ttl
+
+    @property
+    def success_cutoff(self) -> datetime.datetime:
+        return timezone.now() - maybe_timedelta(self.success_ttl)
+
+    @property
+    def nonsuccess_cutoff(self) -> datetime.datetime:
+        return timezone.now() - maybe_timedelta(self.nonsuccess_ttl)
 
     def get_task_names(self):
         qs = self.TaskModel.objects.values('task_name').annotate(name=GroupBy('task_name'))
@@ -129,24 +140,16 @@ def get_task_names(self):
         return task_names
 
     def clean(self):
-        success_cutoff = timezone.now() - maybe_timedelta(self.expires)
-        failed_cutoff = timezone.now() - maybe_timedelta(
-            max(self.expires, settings.FAILED_CELERY_RESULT_EXPIRES),
-        )
-
         for name in self.get_task_names():
-            logger.debug(
-                'Looking for %s task results succeeded before %s or failed before %s',
-                name,
-                success_cutoff,
-                failed_cutoff,
+            success_q = Q(status=states.SUCCESS, date_modified__lt=self.success_cutoff)
+            nonsuccess_q = (
+                ~Q(status=states.SUCCESS)
+                & Q(date_modified__lt=self.nonsuccess_cutoff)
             )
-            success_q = Q(status=states.SUCCESS, date_modified__lt=success_cutoff)
-            failed_q = Q(status=states.FAILED, date_modified__lt=failed_cutoff)
             queryset = (
                 self.TaskModel.objects
                 .filter(task_name=name)
-                .filter(success_q | failed_q)
+                .filter(success_q | nonsuccess_q)
             )
 
             if not queryset.exists():

tests/api/test_generated_endpoints.py

@@ -33,22 +33,17 @@ def test_one(self, client, endpoint, factory, autocreated_count):
         assert _links.get('next') is None
 
     def test_full_page(self, client, endpoint, factory, autocreated_count):
-        for _ in range(10 - autocreated_count):
-            factory()
-
+        factory.create_batch(10 - autocreated_count)
         resp = client.get('/api/v2/{}/'.format(endpoint))
         assert resp.status_code == 200
-
         _json = resp.json()
         assert len(_json['data']) == 10
         _links = _json.get('links', {})
         assert _links.get('prev') is None
         assert _links.get('next') is None
 
     def test_next_page(self, client, endpoint, factory, autocreated_count):
-        for _ in range(20 - autocreated_count):
-            factory()
-
+        factory.create_batch(20 - autocreated_count)
         resp = client.get('/api/v2/{}/'.format(endpoint))
         assert resp.status_code == 200
 

tests/share/test_celery.py

@@ -1,49 +1,72 @@
-import pytest
-import datetime
-
+import contextlib
+from datetime import timedelta
 from unittest import mock
 
+import pytest
 from django.utils import timezone
 
 from share.celery import TaskResultCleaner, CeleryTaskResult
-
 from tests import factories
 
 
-@pytest.mark.usefixtures('nested_django_db')
-class TestResultArchiver:
+@contextlib.contextmanager
+def long_now(new_now=None):
+    _now = new_now or timezone.now()
+    with mock.patch.object(timezone, 'now', return_value=_now):
+        yield _now
 
-    @pytest.fixture(scope='class', autouse=True)
-    def task_result_data(self, class_scoped_django_db):
-        return factories.CeleryTaskResultFactory.create_batch(100)
+
+@pytest.mark.django_db
+class TestResultCleaner:
 
     def test_delete_false(self):
-        trc = TaskResultCleaner(datetime.timedelta(weeks=520), delete=False)
+        factories.CeleryTaskResultFactory.create_batch(10)
+        trc = TaskResultCleaner(timedelta(weeks=520), delete=False)
         assert trc.delete_queryset(CeleryTaskResult.objects.all()) == 0
-        assert CeleryTaskResult.objects.count() != 0
+        assert CeleryTaskResult.objects.count() == 10
 
     def test_delete_queryset(self):
-        trc = TaskResultCleaner(datetime.timedelta(weeks=520))
-        assert trc.delete_queryset(CeleryTaskResult.objects.all()) == 100
+        factories.CeleryTaskResultFactory.create_batch(10)
+        trc = TaskResultCleaner(timedelta(weeks=520))
+        assert trc.delete_queryset(CeleryTaskResult.objects.all()) == 10
         assert CeleryTaskResult.objects.count() == 0
 
-    def test_get_ttl_default(self):
-        trc = TaskResultCleaner(datetime.timedelta(weeks=520))
-        assert ((timezone.now() - datetime.timedelta(weeks=520)) - trc.get_ttl('non-existant-task')) < datetime.timedelta(seconds=2)
-
-    def test_get_ttl(self):
-        trc = TaskResultCleaner(datetime.timedelta(weeks=520))
-        trc.TASK_TTLS['existant-task'] = datetime.timedelta(days=1)
-        assert ((timezone.now() - datetime.timedelta(days=1)) - trc.get_ttl('existant-task')) < datetime.timedelta(seconds=2)
-
-    def test_clean(self):
-        trc = TaskResultCleaner(0, bucket=mock.Mock())
-        factories.CeleryTaskResultFactory.create_batch(100, status='SUCCESS')
-        trc.clean()
-        assert CeleryTaskResult.objects.count() <= 100  # There's an autouse fixture that makes 100
-
-    def test_clean_chunksize(self):
-        trc = TaskResultCleaner(0, bucket=mock.Mock(), chunk_size=1)
-        factories.CeleryTaskResultFactory.create_batch(100, status='SUCCESS')
-        trc.clean()
-        assert CeleryTaskResult.objects.count() <= 100  # There's an autouse fixture that makes 100
+    def test_success_cutoff(self, settings):
+        with long_now() as _now:
+            trc = TaskResultCleaner(timedelta(days=3).total_seconds())
+            _expected = _now - timedelta(days=3)
+            assert trc.success_cutoff == _expected
+
+    def test_nonsuccess_cutoff(self, settings):
+        with long_now() as _now:
+            trc = TaskResultCleaner(
+                success_ttl=timedelta(days=3),
+                nonsuccess_ttl=timedelta(days=5),
+            )
+            assert trc.success_cutoff == _now - timedelta(days=3)
+            assert trc.nonsuccess_cutoff == _now - timedelta(days=5)
+
+    @pytest.mark.parametrize('batch_size', [1, 1111])
+    def test_clean(self, batch_size):
+        with long_now() as _now:
+            with long_now(_now - timedelta(days=7)):
+                # all should be deleted:
+                factories.CeleryTaskResultFactory.create_batch(10, status='SUCCESS')
+                factories.CeleryTaskResultFactory.create_batch(7, status='FAILED')
+            with long_now(_now - timedelta(days=4)):
+                # successes should be deleted:
+                factories.CeleryTaskResultFactory.create_batch(10, status='SUCCESS')
+                factories.CeleryTaskResultFactory.create_batch(7, status='FAILED')
+            # none should be deleted:
+            factories.CeleryTaskResultFactory.create_batch(10, status='SUCCESS')
+            factories.CeleryTaskResultFactory.create_batch(7, status='FAILED')
+            # end setup
+            assert CeleryTaskResult.objects.count() == 51
+            trc = TaskResultCleaner(
+                success_ttl=timedelta(days=3),
+                nonsuccess_ttl=timedelta(days=5),
+                chunk_size=batch_size,
+            )
+            trc.clean()
+            assert CeleryTaskResult.objects.filter(status='SUCCESS').count() == 10
+            assert CeleryTaskResult.objects.exclude(status='SUCCESS').count() == 14