Fix the logic of checking dataflow job state (#34785)

champon1020 · eladkal · web-flow · commit 8fd5ac6530df · 2023-11-10T20:04:30.000+02:00
Co-authored-by: Elad Kalif &lt;45845474+eladkal@users.noreply.github.com&gt;
diff --git a/airflow/providers/google/cloud/hooks/dataflow.py b/airflow/providers/google/cloud/hooks/dataflow.py
@@ -419,7 +419,9 @@ def _check_dataflow_job_state(self, job) -> bool:
                     "JOB_STATE_DRAINED while it is a batch job"
                 )
 
-        if not self._wait_until_finished and current_state == self._expected_terminal_state:
+        if current_state == self._expected_terminal_state:
+            if self._expected_terminal_state == DataflowJobStatus.JOB_STATE_RUNNING:
+                return not self._wait_until_finished
             return True
 
         if current_state in DataflowJobStatus.AWAITING_STATES:
diff --git a/tests/providers/google/cloud/hooks/test_dataflow.py b/tests/providers/google/cloud/hooks/test_dataflow.py
@@ -1417,6 +1417,10 @@ def test_check_dataflow_job_state_wait_until_finished(
     @pytest.mark.parametrize(
         "job_state, wait_until_finished, expected_result",
         [
+            # DONE
+            (DataflowJobStatus.JOB_STATE_DONE, None, True),
+            (DataflowJobStatus.JOB_STATE_DONE, True, True),
+            (DataflowJobStatus.JOB_STATE_DONE, False, True),
             # RUNNING
             (DataflowJobStatus.JOB_STATE_RUNNING, None, False),
             (DataflowJobStatus.JOB_STATE_RUNNING, True, False),

Original file line number	Diff line number	Diff line change
`@@ -419,7 +419,9 @@ def _check_dataflow_job_state(self, job) -> bool:`
`419`	`419`	`"JOB_STATE_DRAINED while it is a batch job"`
`420`	`420`	`)`
`421`	`421`
`422`		`- if not self._wait_until_finished and current_state == self._expected_terminal_state:`
	`422`	`+ if current_state == self._expected_terminal_state:`
	`423`	`+ if self._expected_terminal_state == DataflowJobStatus.JOB_STATE_RUNNING:`
	`424`	`+ return not self._wait_until_finished`
`423`	`425`	`return True`
`424`	`426`
`425`	`427`	`if current_state in DataflowJobStatus.AWAITING_STATES:`