def test_execute_override_project_id(self, mock_hook, mock_uuid): mock_uuid.return_value = self.job_id mock_hook.return_value.project_id = GCP_PROJECT mock_hook.return_value.wait_for_job.return_value = None mock_hook.return_value.submit_job.return_value.reference.job_id = self.job_id op = DataprocSubmitSparkSqlJobOperator( project_id="other-project", task_id=TASK_ID, region=GCP_LOCATION, gcp_conn_id=GCP_CONN_ID, query=self.query, variables=self.variables, impersonation_chain=IMPERSONATION_CHAIN, ) op.execute(context={}) mock_hook.assert_called_once_with( gcp_conn_id=GCP_CONN_ID, impersonation_chain=IMPERSONATION_CHAIN) mock_hook.return_value.submit_job.assert_called_once_with( project_id="other-project", job=self.other_project_job, location=GCP_LOCATION) mock_hook.return_value.wait_for_job.assert_called_once_with( job_id=self.job_id, location=GCP_LOCATION, project_id="other-project")
def test_execute(self, mock_hook, mock_uuid): mock_uuid.return_value = self.job_id mock_hook.return_value.project_id = GCP_PROJECT mock_hook.return_value.wait_for_job.return_value = None mock_hook.return_value.submit_job.return_value.reference.job_id = self.job_id op = DataprocSubmitSparkSqlJobOperator( task_id=TASK_ID, region=GCP_LOCATION, gcp_conn_id=GCP_CONN_ID, query=self.query, variables=self.variables, ) op.execute(context={}) mock_hook.assert_called_once_with(gcp_conn_id=GCP_CONN_ID) mock_hook.return_value.submit_job.assert_called_once_with( project_id=GCP_PROJECT, job=self.job, location=GCP_LOCATION) mock_hook.return_value.wait_for_job.assert_called_once_with( job_id=self.job_id, location=GCP_LOCATION, project_id=GCP_PROJECT)