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