From 72fb5e5383bedf052b5fa8c2af73b631a9be10dc Mon Sep 17 00:00:00 2001
From: Muhammad Hassan Ahmed <33365802+hassan11196@users.noreply.github.com>
Date: Fri, 11 Oct 2024 16:28:25 +0000
Subject: [PATCH] check UserDrainMode before updating AgentDrainMode

check UserDrainMode before updating AgentDrainMode

Add checks before updating AgentDrainMode

Simplify AgentDrainMode change

Simplify AgentDrainMode change and add logging

Fix type while comparing UserDrainMode

use get to prevent KeyError

Fix Indentation in listDiskUsageOverThreshold
---
 src/python/WMCore/Services/ReqMgrAux/ReqMgrAux.py | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/src/python/WMCore/Services/ReqMgrAux/ReqMgrAux.py b/src/python/WMCore/Services/ReqMgrAux/ReqMgrAux.py
index 58e43562db..97e1e421f3 100644
--- a/src/python/WMCore/Services/ReqMgrAux/ReqMgrAux.py
+++ b/src/python/WMCore/Services/ReqMgrAux/ReqMgrAux.py
@@ -410,8 +410,10 @@ def listDiskUsageOverThreshold(config, updateDB):
 
     if updateDB and not t0Flag:
         agentDrainMode = bool(len(overThresholdDisks))
-        if agentConfig and (agentDrainMode != agentConfig["AgentDrainMode"]):
-            reqMgrAux.updateWMAgentConfig(config.Agent.hostName, {"AgentDrainMode": agentDrainMode},
-                                          inPlace=True)
+        if agentConfig and (agentDrainMode != agentConfig.get("AgentDrainMode")):
+            if agentConfig.get("UserDrainMode") is True:
+                logging.info("Agent is already in UserDrainMode: True mode. No need to update.")
+            else:
+                reqMgrAux.updateWMAgentConfig(config.Agent.hostName, {"AgentDrainMode": agentDrainMode}, inPlace=True)
 
     return overThresholdDisks