From 50b9f1f740fa1b30b25b0f5b701f86a01cf05bb7 Mon Sep 17 00:00:00 2001 From: = Date: Thu, 14 Oct 2021 11:17:11 +0000 Subject: [PATCH 1/2] Fix bug in dataflow hook when no jobs are returned --- .../providers/google/cloud/hooks/dataflow.py | 11 +++++++--- .../google/cloud/hooks/test_dataflow.py | 22 +++++++++++++++++++ 2 files changed, 30 insertions(+), 3 deletions(-) diff --git a/airflow/providers/google/cloud/hooks/dataflow.py b/airflow/providers/google/cloud/hooks/dataflow.py index 5562e8ae7b5e7..fcdf20d379452 100644 --- a/airflow/providers/google/cloud/hooks/dataflow.py +++ b/airflow/providers/google/cloud/hooks/dataflow.py @@ -356,10 +356,15 @@ def _fetch_all_jobs(self) -> List[dict]: .jobs() .list(projectId=self._project_number, location=self._job_location) ) - jobs: List[dict] = [] + all_jobs: List[dict] = [] while request is not None: response = request.execute(num_retries=self._num_retries) - jobs.extend(response["jobs"]) + if response is None: + break + jobs = response.get("jobs") + if jobs is None: + break + all_jobs.extend(jobs) request = ( self._dataflow.projects() @@ -367,7 +372,7 @@ def _fetch_all_jobs(self) -> List[dict]: .jobs() .list_next(previous_request=request, previous_response=response) ) - return jobs + return all_jobs def _fetch_jobs_by_prefix_name(self, prefix_name: str) -> List[dict]: jobs = self._fetch_all_jobs() diff --git a/tests/providers/google/cloud/hooks/test_dataflow.py b/tests/providers/google/cloud/hooks/test_dataflow.py index 370e0b3afa37e..fa841d2c1c18a 100644 --- a/tests/providers/google/cloud/hooks/test_dataflow.py +++ b/tests/providers/google/cloud/hooks/test_dataflow.py @@ -1697,6 +1697,28 @@ def test_fetch_list_job_messages_responses(self): ) assert result == ["response_1"] + def test_fetch_all_jobs_when_no_jobs_returned(self): + # fmt: off + mock_list = ( + self.mock_dataflow + .projects.return_value + .locations.return_value + .jobs.return_value + .list + ) + + # fmt: on + mock_list.return_value.execute.return_value = {} + + jobs_controller = _DataflowJobsController( + dataflow=self.mock_dataflow, + project_number=TEST_PROJECT, + location=TEST_LOCATION, + job_id=TEST_JOB_ID, + ) + result = jobs_controller._fetch_all_jobs() + assert result == [] + @mock.patch(DATAFLOW_STRING.format('_DataflowJobsController._fetch_list_job_messages_responses')) def test_fetch_job_messages_by_id(self, mock_fetch_responses): mock_fetch_responses.return_value = iter( From 8b7eaa99dc414428bc59983c2af8435c66461e1f Mon Sep 17 00:00:00 2001 From: keze Date: Thu, 21 Oct 2021 15:52:50 +0000 Subject: [PATCH 2/2] Remove unnecessary check and fix formatting in tests --- airflow/providers/google/cloud/hooks/dataflow.py | 2 -- .../providers/google/cloud/hooks/test_dataflow.py | 15 +++++++-------- 2 files changed, 7 insertions(+), 10 deletions(-) diff --git a/airflow/providers/google/cloud/hooks/dataflow.py b/airflow/providers/google/cloud/hooks/dataflow.py index fcdf20d379452..5063743b76178 100644 --- a/airflow/providers/google/cloud/hooks/dataflow.py +++ b/airflow/providers/google/cloud/hooks/dataflow.py @@ -359,8 +359,6 @@ def _fetch_all_jobs(self) -> List[dict]: all_jobs: List[dict] = [] while request is not None: response = request.execute(num_retries=self._num_retries) - if response is None: - break jobs = response.get("jobs") if jobs is None: break diff --git a/tests/providers/google/cloud/hooks/test_dataflow.py b/tests/providers/google/cloud/hooks/test_dataflow.py index fa841d2c1c18a..f782594de084b 100644 --- a/tests/providers/google/cloud/hooks/test_dataflow.py +++ b/tests/providers/google/cloud/hooks/test_dataflow.py @@ -1699,16 +1699,15 @@ def test_fetch_list_job_messages_responses(self): def test_fetch_all_jobs_when_no_jobs_returned(self): # fmt: off - mock_list = ( + ( self.mock_dataflow - .projects.return_value - .locations.return_value - .jobs.return_value - .list - ) - + .projects.return_value + .locations.return_value + .jobs.return_value + .list.return_value + .execute.return_value + ) = {} # fmt: on - mock_list.return_value.execute.return_value = {} jobs_controller = _DataflowJobsController( dataflow=self.mock_dataflow,