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)
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 } } } } })