Batch: list_jobs() now supports the arrayJobId-parameter (#8321)

getmoto · Nov 16, 2024 · c8d0e1c · c8d0e1c
1 parent 114e344
commit c8d0e1c
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 6 deletions.
diff --git a/moto/batch/models.py b/moto/batch/models.py
@@ -1835,25 +1835,32 @@ def describe_jobs(self, jobs: Optional[List[str]]) -> List[Dict[str, Any]]:
 
     def list_jobs(
         self,
-        job_queue_name: str,
+        job_queue_name: Optional[str],
+        array_job_id: Optional[str],
         job_status: Optional[str] = None,
         filters: Optional[List[Dict[str, Any]]] = None,
     ) -> List[Job]:
         """
         Pagination is not yet implemented
         """
+        jobs_to_check = []
         jobs = []
 
-        job_queue = self.get_job_queue(job_queue_name)
-        if job_queue is None:
-            raise ClientException(f"Job queue {job_queue_name} does not exist")
+        if job_queue_name:
+            if job_queue := self.get_job_queue(job_queue_name):
+                jobs_to_check.extend(job_queue.jobs)
+            else:
+                raise ClientException(f"Job queue {job_queue_name} does not exist")
+        if array_job_id:
+            if array_job := self.get_job_by_id(array_job_id):
+                jobs_to_check.extend(array_job._child_jobs or [])
 
         if job_status is not None and job_status not in JobStatus.job_statuses():
             raise ClientException(
                 "Job status is not one of SUBMITTED | PENDING | RUNNABLE | STARTING | RUNNING | SUCCEEDED | FAILED"
             )
 
-        for job in job_queue.jobs:
+        for job in jobs_to_check:
             if job_status is not None and job.status != job_status:
                 continue
 

diff --git a/moto/batch/responses.py b/moto/batch/responses.py
@@ -215,8 +215,14 @@ def listjobs(self) -> str:
         job_queue = self._get_param("jobQueue")
         job_status = self._get_param("jobStatus")
         filters = self._get_param("filters")
+        array_job_id = self._get_param("arrayJobId")
 
-        jobs = self.batch_backend.list_jobs(job_queue, job_status, filters)
+        jobs = self.batch_backend.list_jobs(
+            job_queue_name=job_queue,
+            array_job_id=array_job_id,
+            job_status=job_status,
+            filters=filters,
+        )
 
         result = {"jobSummaryList": [job.describe_short() for job in jobs]}
         return json.dumps(result)

diff --git a/tests/test_batch/test_batch_jobs.py b/tests/test_batch/test_batch_jobs.py
@@ -130,6 +130,11 @@ def test_submit_job_array_size():
     # Child job was executed
     assert len(child_job_1["attempts"]) == 1
 
+    # List all child jobs
+    child_job_list = batch_client.list_jobs(arrayJobId=job_id)["jobSummaryList"]
+    assert len(child_job_list) == 2
+    assert child_job_1_id in [c["jobId"] for c in child_job_list]
+
 
 @mock_aws
 @pytest.mark.network