Update logic to allow retries in AWS Batch Client hook to be effective (

#38998)

Co-authored-by: evgenyslab <e.nuger@gmail.com>

airflow/providers/amazon/aws/hooks/batch_client.py

@@ -397,6 +397,8 @@ def get_job_description(self, job_id: str) -> dict:
             try:
                 response = self.get_conn().describe_jobs(jobs=[job_id])
                 return self.parse_job_description(job_id, response)
+            except AirflowException as err:
+                self.log.warning(err)
             except botocore.exceptions.ClientError as err:
                 # Allow it to retry in case of exceeded quota limit of requests to AWS API
                 if err.response.get("Error", {}).get("Code") != "TooManyRequestsException":

tests/providers/amazon/aws/hooks/test_batch_client.py

@@ -245,13 +245,14 @@ def test_check_job_success_raises_unknown_status(self):
         assert msg in str(ctx.value)
         assert status in str(ctx.value)
 
-    def test_check_job_success_raises_without_jobs(self):
+    def test_check_job_success_raises_without_jobs(self, caplog):
         self.client_mock.describe_jobs.return_value = {"jobs": []}
-        with pytest.raises(AirflowException) as ctx:
-            self.batch_client.check_job_success(JOB_ID)
-        self.client_mock.describe_jobs.assert_called_once_with(jobs=[JOB_ID])
-        msg = f"AWS Batch job ({JOB_ID}) description error"
-        assert msg in str(ctx.value)
+        with caplog.at_level(level=logging.WARNING):
+            with pytest.raises(AirflowException):
+                self.batch_client.check_job_success(JOB_ID)
+            self.client_mock.describe_jobs.assert_has_calls([mock.call(jobs=[JOB_ID])] * 3)
+            msg = f"AWS Batch job ({JOB_ID}) description error"
+            assert msg in caplog.messages[0]
 
     def test_terminate_job(self):
         self.client_mock.terminate_job.return_value = {}