def test_builder(self, mock_hook, mock_uuid):
mock_hook.return_value.project_id = GCP_PROJECT
mock_uuid.return_value = self.job_id
op = DataprocSubmitPigJobOperator(
task_id=TASK_ID,
region=GCP_LOCATION,
gcp_conn_id=GCP_CONN_ID,
query=self.query,
variables=self.variables,
)
job = op.generate_job()
self.assertDictEqual(self.job, job)
def test_deprecation_warning(self, mock_hook):
with self.assertWarns(DeprecationWarning) as warning:
DataprocSubmitPigJobOperator(
task_id=TASK_ID,
region=GCP_LOCATION,
query="query",
)
assert_warning("DataprocSubmitJobOperator", warning)
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 = DataprocSubmitPigJobOperator(
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)
