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

Co-authored-by: Elad Kalif <[email protected]>
apache · Nov 10, 2023 · 8fd5ac6 · 8fd5ac6
1 parent 3cb0870
commit 8fd5ac6
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 1 deletion.
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),