refactor: use ObjectLoader for DB instantiation in JobCleaningAgent

Author: fstagni

src/DIRAC/WorkloadManagementSystem/Agent/JobCleaningAgent.py

@@ -1,4 +1,4 @@
-""" The Job Cleaning Agent controls removing jobs from the WMS in the end of their life cycle.
+"""The Job Cleaning Agent controls removing jobs from the WMS in the end of their life cycle.
 
     This agent will take care of:
     - removing all jobs that are in status JobStatus.DELETED
@@ -22,6 +22,7 @@
 than 0.
 
 """
+
 import datetime
 import os
 
@@ -30,14 +31,13 @@
 from DIRAC.ConfigurationSystem.Client.Helpers.Registry import getDNForUsername
 from DIRAC.Core.Base.AgentModule import AgentModule
 from DIRAC.Core.Utilities import TimeUtilities
+from DIRAC.Core.Utilities.ObjectLoader import ObjectLoader
 from DIRAC.RequestManagementSystem.Client.File import File
 from DIRAC.RequestManagementSystem.Client.Operation import Operation
 from DIRAC.RequestManagementSystem.Client.ReqClient import ReqClient
 from DIRAC.RequestManagementSystem.Client.Request import Request
 from DIRAC.WorkloadManagementSystem.Client import JobStatus
 from DIRAC.WorkloadManagementSystem.Client.WMSClient import WMSClient
-from DIRAC.WorkloadManagementSystem.DB.JobDB import JobDB
-from DIRAC.WorkloadManagementSystem.DB.SandboxMetadataDB import SandboxMetadataDB
 from DIRAC.WorkloadManagementSystem.DB.StatusUtils import kill_delete_jobs
 from DIRAC.WorkloadManagementSystem.Service.JobPolicy import RIGHT_DELETE
 from DIRAC.WorkloadManagementSystem.Utilities.JobParameters import getJobParameters
@@ -54,6 +54,8 @@ def __init__(self, *args, **kwargs):
 
         # clients
         self.jobDB = None
+        self.taskQueueDB = None
+        self.pilotAgentsDB = None
         self.sandboxDB = None
 
         self.maxJobsAtOnce = 500
@@ -66,8 +68,25 @@ def __init__(self, *args, **kwargs):
     def initialize(self):
         """Sets defaults"""
 
-        self.jobDB = JobDB()
-        self.sandboxDB = SandboxMetadataDB()
+        result = ObjectLoader().loadObject("WorkloadManagementSystem.DB.JobDB", "JobDB")
+        if not result["OK"]:
+            return result
+        self.jobDB = result["Value"]()
+
+        result = ObjectLoader().loadObject("WorkloadManagementSystem.DB.TaskQueueDB", "TaskQueueDB")
+        if not result["OK"]:
+            return result
+        self.taskQueueDB = result["Value"]()
+
+        result = ObjectLoader().loadObject("WorkloadManagementSystem.DB.PilotAgentsDB", "PilotAgentsDB")
+        if not result["OK"]:
+            return result
+        self.pilotAgentsDB = result["Value"]()
+
+        result = ObjectLoader().loadObject("WorkloadManagementSystem.DB.SandboxMetadataDB", "SandboxMetadataDB")
+        if not result["OK"]:
+            return result
+        self.sandboxDB = result["Value"]()
 
         agentTSTypes = self.am_getOption("ProductionTypes", [])
         if agentTSTypes:
@@ -238,7 +257,13 @@ def _deleteRemoveJobs(self, jobList, remove=False):
                 wmsClient = WMSClient(useCertificates=True, delegatedDN=res["Value"][0], delegatedGroup=ownerGroup)
                 result = wmsClient.removeJob(jobsList)
             else:
-                result = kill_delete_jobs(RIGHT_DELETE, jobsList)
+                result = kill_delete_jobs(
+                    RIGHT_DELETE,
+                    jobsList,
+                    jobdb=self.jobDB,
+                    taskqueuedb=self.taskQueueDB,
+                    pilotagentsdb=self.pilotAgentsDB,
+                )
             if not result["OK"]:
                 self.log.error(
                     f"Could not {'remove' if remove else 'delete'} jobs",

src/DIRAC/WorkloadManagementSystem/Agent/test/Test_Agent_JobCleaningAgent.py

@@ -1,5 +1,5 @@
-""" Test class for Job Cleaning Agent
-"""
+"""Test class for Job Cleaning Agent"""
+
 from unittest.mock import MagicMock
 
 import pytest
@@ -16,6 +16,9 @@
 mockNone = MagicMock()
 mockNone.return_value = None
 mockJMC = MagicMock()
+mockJobDB = MagicMock()
+mockJobDB.getDistinctJobAttributes = mockReply
+mockJobDB.selectJobs = mockReply
 
 
 @pytest.fixture
@@ -27,16 +30,20 @@ def jca(mocker):
         create=True,
     )
     mocker.patch("DIRAC.WorkloadManagementSystem.Agent.JobCleaningAgent.AgentModule.am_getOption", return_value=mockAM)
+
+    def mock_load_object(module_path, class_name):
+        mocks = {
+            "JobDB": mockJobDB,
+            "TaskQueueDB": MagicMock(),
+            "PilotAgentsDB": MagicMock(),
+            "SandboxMetadataDB": MagicMock(),
+        }
+        return {"OK": True, "Value": lambda: mocks[class_name]}
+
     mocker.patch(
-        "DIRAC.WorkloadManagementSystem.Agent.JobCleaningAgent.JobDB.getDistinctJobAttributes", side_effect=mockReply
-    )
-    mocker.patch("DIRAC.WorkloadManagementSystem.Agent.JobCleaningAgent.JobDB.selectJobs", side_effect=mockReply)
-    mocker.patch("DIRAC.WorkloadManagementSystem.Agent.JobCleaningAgent.JobDB.__init__", side_effect=mockNone)
-    mocker.patch(
-        "DIRAC.WorkloadManagementSystem.Agent.JobCleaningAgent.SandboxMetadataDB.__init__", side_effect=mockNone
+        "DIRAC.WorkloadManagementSystem.Agent.StalledJobAgent.ObjectLoader.loadObject",
+        side_effect=mock_load_object,
     )
-    mocker.patch("DIRAC.WorkloadManagementSystem.Agent.JobCleaningAgent.ReqClient", return_value=mockNone)
-
     jca = JobCleaningAgent()
     jca.log = gLogger
     jca.log.setLevel("DEBUG")
@@ -128,15 +135,27 @@ def test_deleteJobOversizedSandbox(mocker, inputs, params, expected):
 
     mocker.patch("DIRAC.WorkloadManagementSystem.Agent.JobCleaningAgent.AgentModule.__init__")
     mocker.patch("DIRAC.WorkloadManagementSystem.Agent.JobCleaningAgent.AgentModule.am_getOption", return_value=mockAM)
-    mocker.patch("DIRAC.WorkloadManagementSystem.Agent.JobCleaningAgent.JobDB", return_value=mockNone)
-    mocker.patch("DIRAC.WorkloadManagementSystem.Agent.JobCleaningAgent.SandboxMetadataDB", return_value=mockNone)
     mocker.patch("DIRAC.WorkloadManagementSystem.Agent.JobCleaningAgent.ReqClient", return_value=mockNone)
     mocker.patch(
         "DIRAC.WorkloadManagementSystem.Agent.JobCleaningAgent.getDNForUsername", return_value=S_OK(["/bih/boh/DN"])
     )
     mocker.patch("DIRAC.WorkloadManagementSystem.Agent.JobCleaningAgent.getJobParameters", return_value=params)
 
+    def mock_load_object(module_path, class_name):
+        mocks = {
+            "JobDB": MagicMock(),
+            "TaskQueueDB": MagicMock(),
+            "PilotAgentsDB": MagicMock(),
+            "SandboxMetadataDB": MagicMock(),
+        }
+        return {"OK": True, "Value": lambda: mocks[class_name]}
+
+    mocker.patch(
+        "DIRAC.WorkloadManagementSystem.Agent.StalledJobAgent.ObjectLoader.loadObject",
+        side_effect=mock_load_object,
+    )
     jobCleaningAgent = JobCleaningAgent()
+
     jobCleaningAgent.log = gLogger
     jobCleaningAgent.log.setLevel("DEBUG")
     jobCleaningAgent._AgentModule__configDefaults = mockAM