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.
mock_jobs_list.return_value.execute.return_value = {
            "jobs": [
                {
                    "id": "id-2",
                    "name": "name-2",
                    "currentState": DataflowJobStatus.JOB_STATE_RUNNING,
                    "type": DataflowJobType.JOB_TYPE_STREAMING
                }
            ]
        }
        (
            self.mock_dataflow.projects.return_value.
            locations.return_value.
            jobs.return_value.
            list_next.return_value
        ) = None
        
        dataflow_job = _DataflowJobsController(
            dataflow=self.mock_dataflow,
            project_number=TEST_PROJECT,
            name="name-",
            location=TEST_LOCATION,
            poll_sleep=0,
            job_id=None,
            num_retries=20,
            multiple_jobs=True,
        )
        dataflow_job.wait_for_done()
        assert 1 == mock_jobs_list.call_count
    def test_dataflow_job_wait_for_single_jobs(self):