def test_execute(self, mock_hook):
cluster_update = {
"config": {"worker_config": {"num_instances": 3}, "secondary_worker_config": {"num_instances": 4}}
}
op = DataprocScaleClusterOperator(
task_id=TASK_ID,
cluster_name=CLUSTER_NAME,
project_id=GCP_PROJECT,
region=GCP_LOCATION,
num_workers=3,
num_preemptible_workers=4,
graceful_decommission_timeout="10m",
gcp_conn_id=GCP_CONN_ID,
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.update_cluster.assert_called_once_with(
project_id=GCP_PROJECT,
location=GCP_LOCATION,
cluster_name=CLUSTER_NAME,
cluster=cluster_update,
graceful_decommission_timeout={"seconds": 600},
update_mask=UPDATE_MASK,
)
def test_deprecation_warning(self):
with self.assertWarns(DeprecationWarning) as warning:
DataprocScaleClusterOperator(task_id=TASK_ID, cluster_name=CLUSTER_NAME, project_id=GCP_PROJECT)
assert_warning("DataprocUpdateClusterOperator", warning)