def test_execute_if_cluster_exists(self, mock_hook):
mock_hook.return_value.create_cluster.side_effect = [
AlreadyExists("test")
]
op = DataprocClusterCreateOperator(
task_id=TASK_ID,
region=GCP_LOCATION,
project_id=GCP_PROJECT,
cluster=CLUSTER,
gcp_conn_id=GCP_CONN_ID,
retry=RETRY,
timeout=TIMEOUT,
metadata=METADATA,
request_id=REQUEST_ID,
)
op.execute(context={})
mock_hook.assert_called_once_with(gcp_conn_id=GCP_CONN_ID)
mock_hook.return_value.create_cluster.assert_called_once_with(
region=GCP_LOCATION,
project_id=GCP_PROJECT,
cluster=CLUSTER,
request_id=REQUEST_ID,
retry=RETRY,
timeout=TIMEOUT,
metadata=METADATA,
)
mock_hook.return_value.get_cluster.assert_called_once_with(
region=GCP_LOCATION,
project_id=GCP_PROJECT,
cluster_name=CLUSTER_NAME,
retry=RETRY,
timeout=TIMEOUT,
metadata=METADATA,
)
def test_execute(self, mock_hook):
op = DataprocClusterCreateOperator(
task_id=TASK_ID,
region=GCP_LOCATION,
project_id=GCP_PROJECT,
cluster=CLUSTER,
request_id=REQUEST_ID,
gcp_conn_id=GCP_CONN_ID,
retry=RETRY,
timeout=TIMEOUT,
metadata=METADATA,
)
op.execute(context={})
mock_hook.assert_called_once_with(gcp_conn_id=GCP_CONN_ID)
mock_hook.return_value.create_cluster.assert_called_once_with(
region=GCP_LOCATION,
project_id=GCP_PROJECT,
cluster=CLUSTER,
request_id=REQUEST_ID,
retry=RETRY,
timeout=TIMEOUT,
metadata=METADATA,
)
def test_depreciation_warning(self, mock_generator, mock_signature):
mock_signature.return_value.parameters = cluster_params
with self.assertWarns(DeprecationWarning) as warning:
DataprocClusterCreateOperator(
task_id=TASK_ID,
region=GCP_LOCATION,
project_id=GCP_PROJECT,
cluster_name="cluster_name",
num_workers=2,
zone="zone",
)
assert_warning("Passing cluster parameters by keywords", warning)
mock_generator.assert_called_once_with(
task_id=TASK_ID,
region=GCP_LOCATION,
project_id=GCP_PROJECT,
cluster_name="cluster_name",
num_workers=2,
zone="zone",
)
"cluster_name": CLUSTER_NAME
},
"hadoop_job": {
"main_jar_file_uri":
"file:///usr/lib/hadoop-mapreduce/hadoop-mapreduce-examples.jar",
"args": ["wordcount", "gs://pub/shakespeare/rose.txt", OUTPUT_PATH],
},
}
with models.DAG(
"example_gcp_dataproc",
default_args={"start_date": days_ago(1)},
schedule_interval=None,
) as dag:
create_cluster = DataprocClusterCreateOperator(task_id="create_cluster",
project_id=PROJECT_ID,
cluster=CLUSTER,
region=REGION)
scale_cluster = DataprocUpdateClusterOperator(
task_id="scale_cluster",
cluster_name=CLUSTER_NAME,
cluster=CLUSTER_UPDATE,
update_mask=UPDATE_MASK,
graceful_decommission_timeout=TIMEOUT,
project_id=PROJECT_ID,
location=REGION,
)
pig_task = DataprocSubmitJobOperator(task_id="pig_task",
job=PIG_JOB,
location=REGION,