From b434958a79cd85beaf18be9492e36eed1bec0673 Mon Sep 17 00:00:00 2001
From: Antoine Nguyen <tonio@ngyn.org>
Date: Wed, 26 Mar 2025 14:24:48 +0100
Subject: [PATCH] Make sure to use the right client everywhere

---
 modoboa/policyd/handlers.py | 15 ++-------------
 modoboa/policyd/tests.py    |  8 ++------
 modoboa/policyd/utils.py    | 12 +++++++++---
 3 files changed, 13 insertions(+), 22 deletions(-)

diff --git a/modoboa/policyd/handlers.py b/modoboa/policyd/handlers.py
index 78cb8ef48..04d8e7e1c 100644
--- a/modoboa/policyd/handlers.py
+++ b/modoboa/policyd/handlers.py
@@ -1,14 +1,12 @@
 """App related signal handlers."""
 
-import redis
-
-from django.conf import settings
 from django.db.models import signals
 from django.dispatch import receiver
 
 from modoboa.admin import models as admin_models
 
 from . import constants
+from . import utils
 
 
 def set_message_limit(instance, key):
@@ -16,16 +14,7 @@ def set_message_limit(instance, key):
     old_message_limit = instance._loaded_values.get("message_limit")
     if old_message_limit == instance.message_limit:
         return
-    if not getattr(settings, "REDIS_SENTINEL", False):
-        rclient = redis.Redis(
-            host=settings.REDIS_HOST,
-            port=settings.REDIS_PORT,
-            db=settings.REDIS_QUOTA_DB,
-        )
-    else:
-        sentinel = redis.sentinel.Sentinel(settings.REDIS_SENTINELS, socket_timeout=0.1)
-        rclient = sentinel.master_for(settings.REDIS_MASTER, socket_timeout=0.1)
-
+    rclient = utils.get_redis_connection()
     if instance.message_limit is None:
         # delete existing key
         if rclient.hexists(constants.REDIS_HASHNAME, key):
diff --git a/modoboa/policyd/tests.py b/modoboa/policyd/tests.py
index 000e541f3..0776b9e91 100644
--- a/modoboa/policyd/tests.py
+++ b/modoboa/policyd/tests.py
@@ -22,6 +22,7 @@ from modoboa.lib.tests import ModoTestCase, ParametersMixin
 from modoboa.policyd import core as policyd_core
 
 from . import constants
+from . import utils
 
 
 def start_policy_daemon():
@@ -33,12 +34,7 @@ class RedisTestCaseMixin:
 
     def setUp(self):
         super().setUp()
-        self.rclient = redis.Redis(
-            host=settings.REDIS_HOST,
-            port=settings.REDIS_PORT,
-            db=settings.REDIS_QUOTA_DB,
-        )
-        self.rclient.set_response_callback("HGET", int)
+        self.rclient = utils.get_redis_connection()
         self.rclient.delete(constants.REDIS_HASHNAME)
 
 
diff --git a/modoboa/policyd/utils.py b/modoboa/policyd/utils.py
index ca4b88c17..5c2a06b1b 100644
--- a/modoboa/policyd/utils.py
+++ b/modoboa/policyd/utils.py
@@ -9,9 +9,15 @@ from . import constants
 
 def get_redis_connection():
     """Return a client connection to Redis server."""
-    rclient = redis.Redis(
-        host=settings.REDIS_HOST, port=settings.REDIS_PORT, db=settings.REDIS_QUOTA_DB
-    )
+    if not getattr(settings, "REDIS_SENTINEL", False):
+        rclient = redis.Redis(
+            host=settings.REDIS_HOST,
+            port=settings.REDIS_PORT,
+            db=settings.REDIS_QUOTA_DB,
+        )
+    else:
+        sentinel = redis.sentinel.Sentinel(settings.REDIS_SENTINELS, socket_timeout=0.1)
+        rclient = sentinel.master_for(settings.REDIS_MASTER, socket_timeout=0.1)
     rclient.set_response_callback("HGET", int)
     return rclient