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.
location=DEFAULT_DATAFLOW_LOCATION,
            name=f"test-dataflow-pipeline-{MOCK_UUID_PREFIX}",
            num_retries=5,
            poll_sleep=10,
            project_number=TEST_PROJECT,
            drain_pipeline=False,
            cancel_timeout=DEFAULT_CANCEL_TIMEOUT,
            wait_until_finished=None,
        )
        mock_uuid.assert_called_once_with()
    @mock.patch(DATAFLOW_STRING.format("uuid.uuid4"), return_value=MOCK_UUID)
    @mock.patch(DATAFLOW_STRING.format("_DataflowJobsController"))
    @mock.patch(DATAFLOW_STRING.format("DataflowHook.get_conn"))
    def test_start_template_dataflow_update_runtime_env(self, mock_conn, mock_dataflowjob, mock_uuid):
        options_with_runtime_env = copy.deepcopy(RUNTIME_ENV)
        del options_with_runtime_env["numWorkers"]
        runtime_env = {"numWorkers": 17}
        expected_runtime_env = copy.deepcopy(RUNTIME_ENV)
        expected_runtime_env.update(runtime_env)
        dataflowjob_instance = mock_dataflowjob.return_value
        dataflowjob_instance.wait_for_done.return_value = None
        
        method = (mock_conn.return_value
                  .projects.return_value
                  .locations.return_value
                  .templates.return_value
                  .launch)
        
        method.return_value.execute.return_value = {"job": {"id": TEST_JOB_ID}}
        self.dataflow_hook.start_template_dataflow(