As discussed in  https://github.com/apache/airflow/pull/15534#discussion_r620500075, we could possibly rewrite Beam Hook to remove the need of storing state in it.
project_number=TEST_PROJECT,
            name=UNIQUE_JOB_NAME,
            location=TEST_LOCATION,
            poll_sleep=10,
            job_id=TEST_JOB_ID,
            num_retries=20,
            multiple_jobs=True,
        )
        dataflow_job.wait_for_done()
        
        self.mock_dataflow.projects.return_value.locations.return_value.jobs.return_value.\
            list.assert_called_once_with(location=TEST_LOCATION, projectId=TEST_PROJECT)
        self.mock_dataflow.projects.return_value.locations.return_value.jobs.return_value.list\
            .return_value.execute.assert_called_once_with(num_retries=20)
        
        assert dataflow_job.get_jobs() == [job, job]
    @pytest.mark.parametrize(
        "state, exception_regex",
        [
            (DataflowJobStatus.JOB_STATE_FAILED, "Google Cloud Dataflow job name-2 has failed\\."),
            (DataflowJobStatus.JOB_STATE_CANCELLED, "Google Cloud Dataflow job name-2 was cancelled\\."),
            (DataflowJobStatus.JOB_STATE_DRAINED, "Google Cloud Dataflow job name-2 was drained\\."),
            (DataflowJobStatus.JOB_STATE_UPDATED, "Google Cloud Dataflow job name-2 was updated\\."),
            (
                DataflowJobStatus.JOB_STATE_UNKNOWN,
                "Google Cloud Dataflow job name-2 was unknown state: JOB_STATE_UNKNOWN",
            ),
        ],
    )
    def test_dataflow_job_wait_for_multiple_jobs_and_one_in_terminal_state(self, state, exception_regex):