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_update.assert_called_once_with(
            body={"requestedState": "JOB_STATE_CANCELLED"},
            jobId="test-job-id",
            location=TEST_LOCATION,
            projectId="test-project",
        )
        mock_update.return_value.execute.assert_called_once_with(num_retries=20)
        mock_sleep.assert_has_calls([mock.call(4), mock.call(4), mock.call(4)])
        mock_timeout.assert_called_once_with(
            seconds=10, error_message="Canceling jobs failed due to timeout (10s): test-job-id"
        )
    @pytest.mark.parametrize(
        "drain_pipeline, job_type, requested_state",
        [
            (False, "JOB_TYPE_BATCH", "JOB_STATE_CANCELLED"),
            (False, "JOB_TYPE_STREAMING", "JOB_STATE_CANCELLED"),
            (True, "JOB_TYPE_BATCH", "JOB_STATE_CANCELLED"),
            (True, "JOB_TYPE_STREAMING", "JOB_STATE_DRAINED"),
        ],
    )
    def test_dataflow_job_cancel_or_drain_job(self, drain_pipeline, job_type, requested_state):
        job = {
            "id": TEST_JOB_ID,
            "name": UNIQUE_JOB_NAME,
            "currentState": DataflowJobStatus.JOB_STATE_RUNNING,
            "type": job_type,
        }
        get_method = self.mock_dataflow.projects.return_value.locations.return_value.jobs.return_value.get
        get_method.return_value.execute.return_value = job