From e4f45eba0ae5ef7bf91cd898af7f7b4ebc0b9218 Mon Sep 17 00:00:00 2001
From: Jens Langhammer <jens.langhammer@beryju.org>
Date: Fri, 16 Oct 2020 14:20:24 +0200
Subject: [PATCH] policies/reputation: implement MonitoredTask

---
 passbook/policies/reputation/tasks.py | 25 +++++++++++++++----------
 passbook/policies/reputation/tests.py |  4 ++--
 2 files changed, 17 insertions(+), 12 deletions(-)

diff --git a/passbook/policies/reputation/tasks.py b/passbook/policies/reputation/tasks.py
index 802e3d195..7baa8a083 100644
--- a/passbook/policies/reputation/tasks.py
+++ b/passbook/policies/reputation/tasks.py
@@ -3,6 +3,7 @@ from django.core.cache import cache
 from structlog import get_logger
 
 from passbook.core.models import User
+from passbook.lib.tasks import MonitoredTask, TaskResult, TaskResultStatus
 from passbook.policies.reputation.models import IPReputation, UserReputation
 from passbook.policies.reputation.signals import (
     CACHE_KEY_IP_PREFIX,
@@ -13,27 +14,26 @@ from passbook.root.celery import CELERY_APP
 LOGGER = get_logger()
 
 
-@CELERY_APP.task()
-def save_ip_reputation():
+@CELERY_APP.task(bind=True, base=MonitoredTask)
+def save_ip_reputation(self: MonitoredTask):
     """Save currently cached reputation to database"""
-    keys = cache.keys(CACHE_KEY_IP_PREFIX + "*")
     objects_to_update = []
-    for key in keys:
-        score = cache.get(key)
+    for key, score in cache.get_many(CACHE_KEY_IP_PREFIX + "*").items():
         remote_ip = key.replace(CACHE_KEY_IP_PREFIX, "")
         rep, _ = IPReputation.objects.get_or_create(ip=remote_ip)
         rep.score = score
         objects_to_update.append(rep)
     IPReputation.objects.bulk_update(objects_to_update, ["score"])
+    self.set_status(
+        TaskResult(TaskResultStatus.SUCCESSFUL, ["Successfully updated IP Reputation"])
+    )
 
 
-@CELERY_APP.task()
-def save_user_reputation():
+@CELERY_APP.task(bind=True, base=MonitoredTask)
+def save_user_reputation(self: MonitoredTask):
     """Save currently cached reputation to database"""
-    keys = cache.keys(CACHE_KEY_USER_PREFIX + "*")
     objects_to_update = []
-    for key in keys:
-        score = cache.get(key)
+    for key, score in cache.get_many(CACHE_KEY_USER_PREFIX + "*").items():
         username = key.replace(CACHE_KEY_USER_PREFIX, "")
         users = User.objects.filter(username=username)
         if not users.exists():
@@ -43,3 +43,8 @@ def save_user_reputation():
         rep.score = score
         objects_to_update.append(rep)
     UserReputation.objects.bulk_update(objects_to_update, ["score"])
+    self.set_status(
+        TaskResult(
+            TaskResultStatus.SUCCESSFUL, ["Successfully updated User Reputation"]
+        )
+    )
diff --git a/passbook/policies/reputation/tests.py b/passbook/policies/reputation/tests.py
index 20b643bef..0e2e3d59b 100644
--- a/passbook/policies/reputation/tests.py
+++ b/passbook/policies/reputation/tests.py
@@ -33,7 +33,7 @@ class TestReputationPolicy(TestCase):
         # Test value in cache
         self.assertEqual(cache.get(CACHE_KEY_IP_PREFIX + self.test_ip), -1)
         # Save cache and check db values
-        save_ip_reputation()
+        save_ip_reputation.delay()
         self.assertEqual(IPReputation.objects.get(ip=self.test_ip).score, -1)
 
     def test_user_reputation(self):
@@ -43,7 +43,7 @@ class TestReputationPolicy(TestCase):
         # Test value in cache
         self.assertEqual(cache.get(CACHE_KEY_USER_PREFIX + self.test_username), -1)
         # Save cache and check db values
-        save_user_reputation()
+        save_user_reputation.delay()
         self.assertEqual(UserReputation.objects.get(user=self.user).score, -1)
 
     def test_policy(self):