def test_create_cluster(self): # Setup service.projects().regions().clusters().create() # .execute() self.operation = {'name': 'operation', 'done': True} self.mock_execute = Mock() self.mock_execute.execute.return_value = self.operation self.mock_clusters = Mock() self.mock_clusters.create.return_value = self.mock_execute self.mock_regions = Mock() self.mock_regions.clusters.return_value = self.mock_clusters self.mock_projects = Mock() self.mock_projects.regions.return_value = self.mock_regions self.mock_conn = Mock() self.mock_conn.projects.return_value = self.mock_projects with patch(HOOK) as MockHook: hook = MockHook() hook.get_conn.return_value = self.mock_conn hook.wait.return_value = None dataproc_task = DataprocClusterCreateOperator( task_id=TASK_ID, region=GCP_REGION, cluster_name=CLUSTER_NAME, project_id=GCP_PROJECT_ID, num_workers=NUM_WORKERS, zone=GCE_ZONE, dag=self.dag ) dataproc_task.execute(None) project_uri = 'https://www.googleapis.com/compute/v1/projects/test-project-id' machine_type_uri = project_uri + '/zones/us-central1-a/machineTypes/n1-standard-4' zone_uri = project_uri + '/zones/us-central1-a' self.mock_clusters.create.assert_called_once_with( region=GCP_REGION, projectId=GCP_PROJECT_ID, requestId=mock.ANY, body={ 'projectId': 'test-project-id', 'clusterName': 'test-cluster-name', 'config': { 'gceClusterConfig': {'zoneUri': zone_uri}, 'masterConfig': { 'numInstances': 1, 'machineTypeUri': machine_type_uri, 'diskConfig': {'bootDiskType': 'pd-standard', 'bootDiskSizeGb': 1024}}, 'workerConfig': { 'numInstances': 123, 'machineTypeUri': machine_type_uri, 'diskConfig': {'bootDiskType': 'pd-standard', 'bootDiskSizeGb': 1024}}, 'secondaryWorkerConfig': {}, 'softwareConfig': {}, 'lifecycleConfig': {}, 'encryptionConfig': {}}, 'labels': {'airflow-version': mock.ANY}}) hook.wait.assert_called_once_with(self.operation)
def test_cluster_name_log_no_sub(self): with patch('airflow.contrib.operators.dataproc_operator.DataProcHook') \ as mock_hook, patch('logging.info') as l: dataproc_task = DataprocClusterCreateOperator( task_id=TASK_ID, cluster_name=CLUSTER_NAME, project_id=PROJECT_ID, num_workers=NUM_WORKERS, zone=ZONE, dag=self.dag) with self.assertRaises(TypeError) as _: dataproc_task.execute(None) l.assert_called_with(('Creating cluster: ' + CLUSTER_NAME))
def test_cluster_name_log_no_sub(self): with patch('airflow.contrib.operators.dataproc_operator.DataProcHook') as mock_hook: mock_hook.return_value.get_conn = self.mock_conn dataproc_task = DataprocClusterCreateOperator( task_id=TASK_ID, cluster_name=CLUSTER_NAME, project_id=GCP_PROJECT_ID, num_workers=NUM_WORKERS, zone=GCE_ZONE, 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('Creating cluster: %s', CLUSTER_NAME)
def test_cluster_name_log_no_sub(self): with patch('airflow.contrib.operators.dataproc_operator.DataProcHook') as mock_hook: mock_hook.return_value.get_conn = self.mock_conn dataproc_task = DataprocClusterCreateOperator( task_id=TASK_ID, cluster_name=CLUSTER_NAME, project_id=PROJECT_ID, num_workers=NUM_WORKERS, zone=ZONE, dag=self.dag ) with patch.object(dataproc_task.logger, 'info') as mock_info: with self.assertRaises(TypeError) as _: dataproc_task.execute(None) mock_info.assert_called_with('Creating cluster: %s', CLUSTER_NAME)
def test_cluster_name_log_no_sub(self): with patch('airflow.contrib.operators.dataproc_operator.DataProcHook') \ as mock_hook, patch('logging.info') as l: dataproc_task = DataprocClusterCreateOperator( task_id=TASK_ID, cluster_name=CLUSTER_NAME, project_id=PROJECT_ID, num_workers=NUM_WORKERS, zone=ZONE, dag=self.dag ) with self.assertRaises(TypeError) as _: dataproc_task.execute(None) l.assert_called_with(('Creating cluster: ' + CLUSTER_NAME))
def test_cluster_name_log_sub(self): with patch('airflow.contrib.operators.dataproc_operator.DataProcHook') \ as mock_hook, patch('logging.info') as l: dataproc_task = DataprocClusterCreateOperator( task_id=TASK_ID, cluster_name='smoke-cluster-{{ ts_nodash }}', project_id=PROJECT_ID, num_workers=NUM_WORKERS, zone=ZONE, dag=self.dag ) 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) as _: dataproc_task.execute(None) l.assert_called_with(('Creating cluster: 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 = DataprocClusterCreateOperator( task_id=TASK_ID, cluster_name='smoke-cluster-{{ ts_nodash }}', project_id=PROJECT_ID, num_workers=NUM_WORKERS, zone=ZONE, 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) as _: dataproc_task.execute(None) mock_info.assert_called_with('Creating 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 = DataprocClusterCreateOperator( task_id=TASK_ID, cluster_name='smoke-cluster-{{ ts_nodash }}', project_id=GCP_PROJECT_ID, num_workers=NUM_WORKERS, zone=GCE_ZONE, 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('Creating cluster: %s', u'smoke-cluster-testnodash')