Exemple #1
0
    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")
Exemple #2
0
    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)