From 50b9f1f740fa1b30b25b0f5b701f86a01cf05bb7 Mon Sep 17 00:00:00 2001
From: = <dragankesic@google.com>
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 <dragankesic@google.com>
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,
