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.
"sql",
                "query",
                TEST_SQL_QUERY,
                "--project=test-project",
                "--format=value(job.id)",
                "--job-name=test-sql-job-name",
                "--region=custom-location",
                "--bigquery-project=test-project",
                "--bigquery-dataset=test-dataset",
                "--bigquery-table=beam_output",
                "--bigquery-write-disposition=write-truncate",
            ],
            capture_output=True,
        )
        mock_controller.assert_called_once_with(
            dataflow=mock_get_conn.return_value,
            job_id=TEST_JOB_ID,
            location=TEST_LOCATION,
            poll_sleep=10,
            project_number=TEST_PROJECT,
            num_retries=5,
            drain_pipeline=False,
            wait_until_finished=None,
        )
        mock_controller.return_value.wait_for_done.assert_called_once()
        assert result == test_job
    @mock.patch(DATAFLOW_STRING.format("DataflowHook.get_conn"))
    @mock.patch(DATAFLOW_STRING.format("DataflowHook.provide_authorized_gcloud"))
    @mock.patch(DATAFLOW_STRING.format("subprocess.run"))
    def test_start_sql_job(self, mock_run, mock_provide_authorized_gcloud, mock_get_conn):