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.
job = {
            "id": TEST_JOB_ID,
            "name": UNIQUE_JOB_NAME,
            "type": DataflowJobType.JOB_TYPE_BATCH,
            "currentState": DataflowJobStatus.JOB_STATE_DONE,
        }
        
        (
            self.mock_dataflow.projects.return_value.
            locations.return_value.
            jobs.return_value.
            list.return_value.
            execute.return_value
        ) = {
            "jobs": [job, job]
        }
        (
            self.mock_dataflow.projects.return_value.
            locations.return_value.
            jobs.return_value.
            list_next.return_value
        ) = None
        
        dataflow_job = _DataflowJobsController(
            dataflow=self.mock_dataflow,