events: fix monitored task not removing state (#6386)

when `save_on_success` is set, a task failure saves state. when it succeeds afterwards, that state should be removed

Signed-off-by: Jens Langhammer <jens@goauthentik.io>
This commit is contained in:
Jens L 2023-07-26 16:00:50 +02:00 committed by GitHub
parent 17fe595528
commit f272d14fcf
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 56 additions and 6 deletions

View File

@ -76,9 +76,20 @@ class TaskInfo:
return cache.get_many(cache.keys(CACHE_KEY_PREFIX + name)).values()
return cache.get(CACHE_KEY_PREFIX + name, None)
@property
def full_name(self) -> str:
"""Get the full cache key with task name and UID"""
key = CACHE_KEY_PREFIX + self.task_name
if self.result.uid:
uid_suffix = f":{self.result.uid}"
key += uid_suffix
if not self.task_name.endswith(uid_suffix):
self.task_name += uid_suffix
return key
def delete(self):
"""Delete task info from cache"""
return cache.delete(CACHE_KEY_PREFIX + self.task_name)
return cache.delete(self.full_name)
def update_metrics(self):
"""Update prometheus metrics"""
@ -97,12 +108,8 @@ class TaskInfo:
def save(self, timeout_hours=6):
"""Save task into cache"""
key = CACHE_KEY_PREFIX + self.task_name
if self.result.uid:
key += f":{self.result.uid}"
self.task_name += f":{self.result.uid}"
self.update_metrics()
cache.set(key, self, timeout=timeout_hours * 60 * 60)
cache.set(self.full_name, self, timeout=timeout_hours * 60 * 60)
class MonitoredTask(Task):

View File

@ -0,0 +1,43 @@
"""Test Monitored tasks"""
from django.test import TestCase
from authentik.events.monitored_tasks import MonitoredTask, TaskInfo, TaskResult, TaskResultStatus
from authentik.lib.generators import generate_id
from authentik.root.celery import CELERY_APP
class TestMonitoredTasks(TestCase):
"""Test Monitored tasks"""
def test_failed_successful_remove_state(self):
"""Test that a task with `save_on_success` set to `False` that failed saves
a state, and upon successful completion will delete the state"""
should_fail = True
uid = generate_id()
@CELERY_APP.task(
bind=True,
base=MonitoredTask,
)
def test_task(self: MonitoredTask):
self.save_on_success = False
self.set_uid(uid)
self.set_status(
TaskResult(TaskResultStatus.ERROR if should_fail else TaskResultStatus.SUCCESSFUL)
)
# First test successful run
should_fail = False
test_task.delay().get()
self.assertIsNone(TaskInfo.by_name(f"test_task:{uid}"))
# Then test failed
should_fail = True
test_task.delay().get()
info = TaskInfo.by_name(f"test_task:{uid}")
self.assertEqual(info.result.status, TaskResultStatus.ERROR)
# Then after that, the state should be removed
should_fail = False
test_task.delay().get()
self.assertIsNone(TaskInfo.by_name(f"test_task:{uid}"))