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_logging.info = MagicMock()
        mock_logging.warning = MagicMock()
        mock_proc = MagicMock()
        mock_proc.stderr = MagicMock()
        mock_proc.stderr.readlines = MagicMock(return_value=["test\n", "error\n"])
        mock_stderr_fd = MagicMock()
        mock_proc.stderr.fileno = MagicMock(return_value=mock_stderr_fd)
        mock_proc_poll = MagicMock()
        mock_select.return_value = [[mock_stderr_fd]]
        def poll_resp_error():
            mock_proc.return_code = 1
            return True
        mock_proc_poll.side_effect = [None, poll_resp_error]
        mock_proc.poll = mock_proc_poll
        mock_popen.return_value = mock_proc
        dataflow = BeamCommandRunner(["test", "cmd"])
        mock_logging.info.assert_called_once_with("Running command: %s", "test cmd")
        with pytest.raises(Exception):
            dataflow.wait_for_done()
@pytest.fixture()
def make_mock_awaitable():
    def func(mock_obj, return_value):
        f = Future()
        f.set_result(return_value)
        mock_obj.return_value = f
    return func
class TestAsyncHook: