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.
expected_variables["region"] = TEST_LOCATION
        expected_variables["labels"] = '{"foo":"bar"}'
        mock_callback_on_job_id.assert_called_once_with(on_new_job_id_callback)
        mock_beam_start_java_pipeline.assert_called_once_with(
            variables=expected_variables,
            jar=JAR_FILE,
            job_class=JOB_CLASS,
            process_line_callback=mock_callback_on_job_id.return_value,
        )
        mock_dataflow_wait_for_done.assert_called_once_with(
            job_id=mock.ANY, job_name=job_name, location=TEST_LOCATION, multiple_jobs=False
        )
    @mock.patch(DATAFLOW_STRING.format("uuid.uuid4"))
    @mock.patch(DATAFLOW_STRING.format("DataflowHook.wait_for_done"))
    @mock.patch(DATAFLOW_STRING.format("process_line_and_extract_dataflow_job_id_callback"))
    def test_start_java_dataflow_with_custom_region_as_parameter(
        self, mock_callback_on_job_id, mock_dataflow_wait_for_done, mock_uuid
    ):
        mock_beam_start_java_pipeline = self.dataflow_hook.beam_hook.start_java_pipeline
        mock_uuid.return_value = MOCK_UUID
        on_new_job_id_callback = MagicMock()
        job_name = f"{JOB_NAME}-{MOCK_UUID_PREFIX}"
        with pytest.warns(DeprecationWarning, match="This method is deprecated"):
            self.dataflow_hook.start_java_dataflow(
                job_name=JOB_NAME,
                variables=DATAFLOW_VARIABLES_JAVA,
                jar=JAR_FILE,
                job_class=JOB_CLASS,
                on_new_job_id_callback=on_new_job_id_callback,
                location=TEST_LOCATION,