def test_iniline_workflow(self):
with patch(HOOK) as mock_hook:
hook = mock_hook()
hook.get_conn.return_value = self.mock_conn
hook.wait.return_value = None
template = {
"placement": {
"managed_cluster": {
"cluster_name": CLUSTER_NAME,
"config": {
"gce_cluster_config": {
"zone_uri": GCE_ZONE,
}
}
}
},
"jobs": [
{
"step_id": "say-hello",
"pig_job": {
"query": "sh echo hello"
}
}],
}
dataproc_task = DataprocWorkflowTemplateInstantiateInlineOperator(
task_id=TASK_ID,
project_id=GCP_PROJECT_ID,
region=GCP_REGION,
template=template,
dag=self.dag
)
dataproc_task.execute(None)
self.mock_workflows.instantiateInline.assert_called_once_with(
parent='projects/test-project-id/regions/test-region',
requestId=mock.ANY,
body=template)
hook.wait.assert_called_once_with(self.operation)
def test_iniline_workflow(self):
with patch(HOOK) as MockHook:
hook = MockHook()
hook.get_conn.return_value = self.mock_conn
hook.wait.return_value = None
template = {
"placement": {
"managed_cluster": {
"cluster_name": CLUSTER_NAME,
"config": {
"gce_cluster_config": {
"zone_uri": GCE_ZONE,
}
}
}
},
"jobs": [
{
"step_id": "say-hello",
"pig_job": {
"query": "sh echo hello"
}
}],
}
dataproc_task = DataprocWorkflowTemplateInstantiateInlineOperator(
task_id=TASK_ID,
project_id=GCP_PROJECT_ID,
region=GCP_REGION,
template=template,
dag=self.dag
)
dataproc_task.execute(None)
self.mock_workflows.instantiateInline.assert_called_once_with(
parent='projects/test-project-id/regions/test-region',
instanceId=mock.ANY,
body=template)
hook.wait.assert_called_once_with(self.operation)
task_2 = DataprocWorkflowTemplateInstantiateInlineOperator(
task_id='RunWorkflow',
project_id=PROJECT_ID,
region=REGION_ID,
template={
"jobs": [{
"sparkJob": {
"mainJarFileUri": str(SPARK_JAR),
"mainClass": str(SPARK_MAIN_CLASS),
"args": []
},
"stepId": "1_run-job"
}],
"placement": {
"managedCluster": {
"clusterName": "two-node-cluster",
"config": {
"gceClusterConfig": {
"zoneUri": "us-central1-f"
},
"masterConfig": {
"diskConfig": {
"bootDiskSizeGb": 250
},
"machineTypeUri": "n1-standard-2"
},
"softwareConfig": {
"imageVersion": "1.4-deb9"
},
"workerConfig": {
"diskConfig": {
"bootDiskSizeGb": 250
},
"machineTypeUri": "n1-standard-2",
"numInstances": 2
}
}
}
}
})
