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.
py_requirements=None,
            py_system_site_packages=False,
            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=DEFAULT_DATAFLOW_LOCATION
        )
    @pytest.mark.parametrize(
        "current_py_requirements, current_py_system_site_packages",
        [
            (["foo-bar"], False),
            (["foo-bar"], True),
            ([], True),
        ],
    )
    @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_python_dataflow_with_non_empty_py_requirements_and_without_system_packages(
        self,
        mock_callback_on_job_id,
        mock_dataflow_wait_for_done,
        mock_uuid,
        current_py_requirements,
        current_py_system_site_packages,
    ):
        mock_beam_start_python_pipeline = self.dataflow_hook.beam_hook.start_python_pipeline
        mock_uuid.return_value = MOCK_UUID
        on_new_job_id_callback = MagicMock()
        job_name = f"{JOB_NAME}-{MOCK_UUID_PREFIX}"