def test_update_cluster(self):
with patch(HOOK) as mock_hook:
hook = mock_hook()
hook.get_conn.return_value = self.mock_conn
hook.wait.return_value = None
dataproc_task = DataprocClusterScaleOperator(
task_id=TASK_ID,
region=GCP_REGION,
project_id=GCP_PROJECT_ID,
cluster_name=CLUSTER_NAME,
num_workers=NUM_WORKERS,
num_preemptible_workers=NUM_PREEMPTIBLE_WORKERS,
dag=self.dag
)
dataproc_task.execute(None)
self.mock_clusters.patch.assert_called_once_with(
region=GCP_REGION,
projectId=GCP_PROJECT_ID,
clusterName=CLUSTER_NAME,
requestId=mock.ANY,
updateMask="config.worker_config.num_instances,"
"config.secondary_worker_config.num_instances",
body={
'config': {
'workerConfig': {
'numInstances': NUM_WORKERS
},
'secondaryWorkerConfig': {
'numInstances': NUM_PREEMPTIBLE_WORKERS
}
}
})
hook.wait.assert_called_once_with(self.operation)
def test_cluster_name_log_no_sub(self):
with patch('airflow.contrib.hooks.gcp_dataproc_hook.DataProcHook') as mock_hook:
mock_hook.return_value.get_conn = self.mock_conn
dataproc_task = DataprocClusterScaleOperator(
task_id=TASK_ID,
cluster_name=CLUSTER_NAME,
project_id=GCP_PROJECT_ID,
num_workers=NUM_WORKERS,
num_preemptible_workers=NUM_PREEMPTIBLE_WORKERS,
dag=self.dag
)
with patch.object(dataproc_task.log, 'info') as mock_info:
with self.assertRaises(TypeError):
dataproc_task.execute(None)
mock_info.assert_called_with('Scaling cluster: %s', CLUSTER_NAME)
def test_cluster_name_log_no_sub(self):
with patch('airflow.contrib.hooks.gcp_dataproc_hook.DataProcHook') as mock_hook:
mock_hook.return_value.get_conn = self.mock_conn
dataproc_task = DataprocClusterScaleOperator(
task_id=TASK_ID,
cluster_name=CLUSTER_NAME,
project_id=PROJECT_ID,
num_workers=NUM_WORKERS,
num_preemptible_workers=NUM_PREEMPTIBLE_WORKERS,
dag=self.dag
)
with patch.object(dataproc_task.log, 'info') as mock_info:
with self.assertRaises(TypeError):
dataproc_task.execute(None)
mock_info.assert_called_with('Scaling cluster: %s', CLUSTER_NAME)
def test_cluster_name_log_sub(self):
with patch('airflow.contrib.operators.dataproc_operator.DataProcHook') as mock_hook:
mock_hook.return_value.get_conn = self.mock_conn
dataproc_task = DataprocClusterScaleOperator(
task_id=TASK_ID,
cluster_name='smoke-cluster-{{ ts_nodash }}',
project_id=GCP_PROJECT_ID,
num_workers=NUM_WORKERS,
num_preemptible_workers=NUM_PREEMPTIBLE_WORKERS,
dag=self.dag
)
with patch.object(dataproc_task.log, 'info') as mock_info:
context = {'ts_nodash': 'testnodash'}
rendered = dataproc_task.render_template(
'cluster_name',
getattr(dataproc_task, 'cluster_name'), context)
setattr(dataproc_task, 'cluster_name', rendered)
with self.assertRaises(TypeError):
dataproc_task.execute(None)
mock_info.assert_called_with('Scaling cluster: %s', u'smoke-cluster-testnodash')
def test_cluster_name_log_sub(self):
with patch('airflow.contrib.operators.dataproc_operator.DataProcHook') \
as mock_hook:
mock_hook.return_value.get_conn = self.mock_conn
dataproc_task = DataprocClusterScaleOperator(
task_id=TASK_ID,
cluster_name='smoke-cluster-{{ ts_nodash }}',
project_id=PROJECT_ID,
num_workers=NUM_WORKERS,
num_preemptible_workers=NUM_PREEMPTIBLE_WORKERS,
dag=self.dag)
with patch.object(dataproc_task.log, 'info') as mock_info:
context = {'ts_nodash': 'testnodash'}
rendered = dataproc_task.render_template(
'cluster_name', getattr(dataproc_task, 'cluster_name'),
context)
setattr(dataproc_task, 'cluster_name', rendered)
with self.assertRaises(TypeError):
dataproc_task.execute(None)
mock_info.assert_called_with('Scaling cluster: %s',
u'smoke-cluster-testnodash')