diff --git a/allianceauth/services/tasks.py b/allianceauth/services/tasks.py index 6db9d6e8..a53e8693 100644 --- a/allianceauth/services/tasks.py +++ b/allianceauth/services/tasks.py @@ -5,7 +5,6 @@ from django.contrib.auth.models import User from .hooks import ServicesHook from celery_once import QueueOnce as BaseTask, AlreadyQueued from django.core.cache import cache -from time import time logger = logging.getLogger(__name__) @@ -22,14 +21,9 @@ class DjangoBackend: @staticmethod def raise_or_lock(key, timeout): - now = int(time()) - result = cache.get(key) - if result: - remaining = int(result) - now - if remaining > 0: - raise AlreadyQueued(remaining) - else: - cache.set(key, now + timeout, timeout) + acquired = cache.add(key=key, value="lock", timeout=timeout) + if not acquired: + raise AlreadyQueued(int(cache.ttl(key))) @staticmethod def clear_lock(key): diff --git a/allianceauth/services/tests/test_tasks.py b/allianceauth/services/tests/test_tasks.py index 39474690..fae76402 100644 --- a/allianceauth/services/tests/test_tasks.py +++ b/allianceauth/services/tests/test_tasks.py @@ -1,11 +1,15 @@ from unittest import mock +from celery_once import AlreadyQueued + +from django.core.cache import cache from django.test import TestCase from allianceauth.tests.auth_utils import AuthUtils - from allianceauth.services.tasks import validate_services +from ..tasks import DjangoBackend + class ServicesTasksTestCase(TestCase): def setUp(self): @@ -24,3 +28,46 @@ class ServicesTasksTestCase(TestCase): self.assertTrue(svc.validate_user.called) args, kwargs = svc.validate_user.call_args self.assertEqual(self.member, args[0]) # Assert correct user is passed to service hook function + + +class TestDjangoBackend(TestCase): + + TEST_KEY = "my-django-backend-test-key" + TIMEOUT = 1800 + + def setUp(self) -> None: + cache.delete(self.TEST_KEY) + self.backend = DjangoBackend(dict()) + + def test_can_get_lock(self): + """ + when lock can be acquired + then set it with timetout + """ + self.backend.raise_or_lock(self.TEST_KEY, self.TIMEOUT) + self.assertIsNotNone(cache.get(self.TEST_KEY)) + self.assertAlmostEqual(cache.ttl(self.TEST_KEY), self.TIMEOUT, delta=2) + + def test_when_cant_get_lock_raise_exception(self): + """ + when lock can bot be acquired + then raise AlreadyQueued exception with countdown + """ + self.backend.raise_or_lock(self.TEST_KEY, self.TIMEOUT) + + try: + self.backend.raise_or_lock(self.TEST_KEY, self.TIMEOUT) + except Exception as ex: + self.assertIsInstance(ex, AlreadyQueued) + self.assertAlmostEqual(ex.countdown, self.TIMEOUT, delta=2) + + def test_can_clear_lock(self): + """ + when a lock exists + then can get a new lock after clearing it + """ + self.backend.raise_or_lock(self.TEST_KEY, self.TIMEOUT) + + self.backend.clear_lock(self.TEST_KEY) + self.backend.raise_or_lock(self.TEST_KEY, self.TIMEOUT) + self.assertIsNotNone(cache.get(self.TEST_KEY))