fix: use directly the DB instead of going through the service

Author: fstagni

src/DIRAC/WorkloadManagementSystem/Utilities/jobAdministration.py

@@ -1,11 +1,10 @@
 from DIRAC import S_ERROR, S_OK, gLogger
 from DIRAC.StorageManagementSystem.DB.StorageManagementDB import StorageManagementDB
 from DIRAC.WorkloadManagementSystem.Client import JobStatus
-from DIRAC.WorkloadManagementSystem.Client.JobStatus import filterJobStateTransition
 from DIRAC.WorkloadManagementSystem.DB.JobDB import JobDB
 from DIRAC.WorkloadManagementSystem.DB.PilotAgentsDB import PilotAgentsDB
 from DIRAC.WorkloadManagementSystem.DB.TaskQueueDB import TaskQueueDB
-from DIRAC.WorkloadManagementSystem.Service.JobPolicy import RIGHT_KILL, RIGHT_DELETE
+from DIRAC.WorkloadManagementSystem.Service.JobPolicy import RIGHT_DELETE, RIGHT_KILL
 
 
 def _deleteJob(jobID, force=False):
@@ -79,15 +78,20 @@ def kill_delete_jobs(right, validJobList, nonauthJobList=[], force=False):
     killJobList = []
     deleteJobList = []
     if validJobList:
+        result = JobDB().getJobsAttributes(killJobList, ["Status"])
+        if not result["OK"]:
+            return result
+        jobStates = result["Value"]
+
         # Get the jobs allowed to transition to the Killed state
-        filterRes = filterJobStateTransition(validJobList, JobStatus.KILLED)
+        filterRes = _filterJobStateTransition(jobStates, JobStatus.KILLED)
         if not filterRes["OK"]:
             return filterRes
         killJobList.extend(filterRes["Value"])
 
         if right == RIGHT_DELETE:
             # Get the jobs allowed to transition to the Deleted state
-            filterRes = filterJobStateTransition(validJobList, JobStatus.DELETED)
+            filterRes = _filterJobStateTransition(jobStates, JobStatus.DELETED)
             if not filterRes["OK"]:
                 return filterRes
             deleteJobList.extend(filterRes["Value"])
@@ -103,10 +107,7 @@ def kill_delete_jobs(right, validJobList, nonauthJobList=[], force=False):
                 badIDs.append(jobID)
 
         # Look for jobs that are in the Staging state to send kill signal to the stager
-        result = JobDB().getJobsAttributes(killJobList, ["Status"])
-        if not result["OK"]:
-            return result
-        stagingJobList = [jobID for jobID, sDict in result["Value"].items() if sDict["Status"] == JobStatus.STAGING]
+        stagingJobList = [jobID for jobID, sDict in jobStates.items() if sDict["Status"] == JobStatus.STAGING]
 
         if stagingJobList:
             stagerDB = StorageManagementDB()
@@ -127,3 +128,17 @@ def kill_delete_jobs(right, validJobList, nonauthJobList=[], force=False):
 
     jobsList = killJobList if right == RIGHT_KILL else deleteJobList
     return S_OK(jobsList)
+
+
+def _filterJobStateTransition(jobStates, candidateState):
+    """Given a dictionary of jobs states,
+    return a list of jobs that are allowed to transition to the given candidate state.
+    """
+    allowedJobs = []
+
+    for js in jobStates.items():
+        stateRes = JobStatus.JobsStateMachine(js[1]["Status"]).getNextState(candidateState)
+        if stateRes["OK"]:
+            if stateRes["Value"] == candidateState:
+                allowedJobs.append(js[0])
+    return S_OK(allowedJobs)

src/DIRAC/WorkloadManagementSystem/Utilities/test/Test_JobAdministration.py

@@ -10,31 +10,19 @@
 
 
 @pytest.mark.parametrize(
-    "jobIDs_list, right, filtered_jobs, expected_res, expected_value",
+    "jobIDs_list, right",
     [
-        ([], "Kill", [], True, []),
-        ([], "Delete", [], True, []),
-        (1, "Kill", [], True, []),
-        (1, "Kill", [1], True, [1]),
-        ([1, 2], "Kill", [], True, []),
-        ([1, 2], "Kill", [1], True, [1]),
-        (1, "Kill", [1], True, [1]),
-        ([1, 2], "Kill", [1], True, [1]),
-        ([1, 2], "Kill", [2], True, [2]),
-        ([1, 2], "Kill", [], True, []),
-        ([1, 2], "Kill", [1, 2], True, [1, 2]),
+        ([], "Kill"),
+        ([], "Delete"),
+        (1, "Kill"),
+        ([1, 2], "Kill"),
     ],
 )
-def test___kill_delete_jobs(mocker, jobIDs_list, right, filtered_jobs, expected_res, expected_value):
+def test___kill_delete_jobs(mocker, jobIDs_list, right):
     mocker.patch("DIRAC.WorkloadManagementSystem.Utilities.jobAdministration.JobDB", MagicMock())
     mocker.patch("DIRAC.WorkloadManagementSystem.Utilities.jobAdministration.TaskQueueDB", MagicMock())
     mocker.patch("DIRAC.WorkloadManagementSystem.Utilities.jobAdministration.PilotAgentsDB", MagicMock())
     mocker.patch("DIRAC.WorkloadManagementSystem.Utilities.jobAdministration.StorageManagementDB", MagicMock())
-    mocker.patch(
-        "DIRAC.WorkloadManagementSystem.Utilities.jobAdministration.filterJobStateTransition",
-        return_value={"OK": True, "Value": filtered_jobs},
-    )
 
     res = kill_delete_jobs(right, jobIDs_list)
-    assert res["OK"] == expected_res
-    assert res["Value"] == expected_value
+    assert res["OK"]