def execute(self, context):
hook = ComputeEngineHook(gcp_conn_id=self.gcp_conn_id,
api_version=self.api_version)
old_instance_group_manager = hook.get_instance_group_manager(
zone=self.zone,
resource_id=self.resource_id,
project_id=self.project_id)
patch_body = {}
if 'versions' in old_instance_group_manager:
patch_body['versions'] = old_instance_group_manager['versions']
if 'instanceTemplate' in old_instance_group_manager:
patch_body['instanceTemplate'] = old_instance_group_manager[
'instanceTemplate']
if self.update_policy:
patch_body['updatePolicy'] = self.update_policy
self._possibly_replace_template(patch_body)
if 'versions' in patch_body:
for version in patch_body['versions']:
self._possibly_replace_template(version)
if self._change_performed or self.update_policy:
self.log.info(
"Calling patch instance template with updated body: %s",
patch_body)
return hook.patch_instance_group_manager(
zone=self.zone,
resource_id=self.resource_id,
body=patch_body,
request_id=self.request_id,
project_id=self.project_id)
else:
# Idempotence achieved
return True
def execute(self, context):
hook = ComputeEngineHook(gcp_conn_id=self.gcp_conn_id,
api_version=self.api_version)
self._validate_all_body_fields()
return hook.set_machine_type(zone=self.zone,
resource_id=self.resource_id,
body=self.body,
project_id=self.project_id)
def execute(self, context):
hook = ComputeEngineHook(gcp_conn_id=self.gcp_conn_id,
api_version=self.api_version)
self._validate_all_body_fields()
try:
# Idempotence check (sort of) - we want to check if the new template
# is already created and if is, then we assume it was created by previous run
# of CopyTemplate operator - we do not check if content of the template
# is as expected. Templates are immutable so we cannot update it anyway
# and deleting/recreating is not worth the hassle especially
# that we cannot delete template if it is already used in some Instance
# Group Manager. We assume success if the template is simply present
existing_template = hook.get_instance_template(
resource_id=self.body_patch['name'],
project_id=self.project_id)
self.log.info(
"The %s template already existed. It was likely created by previous run of the operator. "
"Assuming success.", existing_template)
return existing_template
except HttpError as e:
# We actually expect to get 404 / Not Found here as the template should
# not yet exist
if not e.resp.status == 404:
raise e
old_body = hook.get_instance_template(resource_id=self.resource_id,
project_id=self.project_id)
new_body = deepcopy(old_body)
self._field_sanitizer.sanitize(new_body)
new_body = merge(new_body, self.body_patch)
self.log.info("Calling insert instance template with updated body: %s",
new_body)
hook.insert_instance_template(body=new_body,
request_id=self.request_id,
project_id=self.project_id)
return hook.get_instance_template(resource_id=self.body_patch['name'],
project_id=self.project_id)
def execute(self, context):
hook = ComputeEngineHook(gcp_conn_id=self.gcp_conn_id,
api_version=self.api_version)
hook.stop_instance(zone=self.zone,
resource_id=self.resource_id,
project_id=self.project_id)
def setUp(self):
with mock.patch('airflow.gcp.hooks.base.GoogleCloudBaseHook.__init__',
new=mock_base_gcp_hook_no_default_project_id):
self.gce_hook_no_project_id = ComputeEngineHook(gcp_conn_id='test')
class TestGcpComputeHookNoDefaultProjectId(unittest.TestCase):
def setUp(self):
with mock.patch('airflow.gcp.hooks.base.GoogleCloudBaseHook.__init__',
new=mock_base_gcp_hook_no_default_project_id):
self.gce_hook_no_project_id = ComputeEngineHook(gcp_conn_id='test')
@mock.patch("airflow.gcp.hooks.compute.ComputeEngineHook._authorize")
@mock.patch("airflow.gcp.hooks.compute.build")
def test_gce_client_creation(self, mock_build, mock_authorize):
result = self.gce_hook_no_project_id.get_conn()
mock_build.assert_called_once_with('compute',
'v1',
http=mock_authorize.return_value,
cache_discovery=False)
self.assertEqual(mock_build.return_value, result)
@mock.patch('airflow.gcp.hooks.compute.ComputeEngineHook.get_conn')
@mock.patch(
'airflow.gcp.hooks.compute.ComputeEngineHook._wait_for_operation_to_complete'
)
def test_start_instance_overridden_project_id(
self, wait_for_operation_to_complete, get_conn):
start_method = get_conn.return_value.instances.return_value.start
execute_method = start_method.return_value.execute
execute_method.return_value = {"name": "operation_id"}
wait_for_operation_to_complete.return_value = None
res = self.gce_hook_no_project_id.start_instance(
project_id='example-project',
zone=GCE_ZONE,
resource_id=GCE_INSTANCE)
self.assertIsNone(res)
start_method.assert_called_once_with(instance='instance',
project='example-project',
zone='zone')
execute_method.assert_called_once_with(num_retries=5)
wait_for_operation_to_complete.assert_called_once_with(
project_id='example-project',
operation_name='operation_id',
zone='zone')
@mock.patch('airflow.gcp.hooks.base.GoogleCloudBaseHook.project_id',
new_callable=PropertyMock,
return_value=None)
@mock.patch('airflow.gcp.hooks.compute.ComputeEngineHook.get_conn')
@mock.patch(
'airflow.gcp.hooks.compute.ComputeEngineHook._wait_for_operation_to_complete'
)
def test_start_instance_no_project_id(self, wait_for_operation_to_complete,
get_conn, mock_project_id):
start_method = get_conn.return_value.instances.return_value.start
execute_method = start_method.return_value.execute
execute_method.return_value = {"name": "operation_id"}
wait_for_operation_to_complete.return_value = None
with self.assertRaises(AirflowException) as cm:
self.gce_hook_no_project_id.start_instance(
zone=GCE_ZONE, resource_id=GCE_INSTANCE)
start_method.assert_not_called()
execute_method.assert_not_called()
err = cm.exception
self.assertIn("The project id must be passed", str(err))
wait_for_operation_to_complete.assert_not_called()
@mock.patch('airflow.gcp.hooks.compute.ComputeEngineHook.get_conn')
@mock.patch(
'airflow.gcp.hooks.compute.ComputeEngineHook._wait_for_operation_to_complete'
)
def test_stop_instance_overridden_project_id(
self, wait_for_operation_to_complete, get_conn):
stop_method = get_conn.return_value.instances.return_value.stop
execute_method = stop_method.return_value.execute
execute_method.return_value = {"name": "operation_id"}
wait_for_operation_to_complete.return_value = None
res = self.gce_hook_no_project_id.stop_instance(
project_id='example-project',
zone=GCE_ZONE,
resource_id=GCE_INSTANCE)
self.assertIsNone(res)
stop_method.assert_called_once_with(instance='instance',
project='example-project',
zone='zone')
execute_method.assert_called_once_with(num_retries=5)
wait_for_operation_to_complete.assert_called_once_with(
project_id='example-project',
operation_name='operation_id',
zone='zone')
@mock.patch('airflow.gcp.hooks.base.GoogleCloudBaseHook.project_id',
new_callable=PropertyMock,
return_value=None)
@mock.patch('airflow.gcp.hooks.compute.ComputeEngineHook.get_conn')
@mock.patch(
'airflow.gcp.hooks.compute.ComputeEngineHook._wait_for_operation_to_complete'
)
def test_stop_instance_no_project_id(self, wait_for_operation_to_complete,
get_conn, mock_project_id):
stop_method = get_conn.return_value.instances.return_value.stop
execute_method = stop_method.return_value.execute
execute_method.return_value = {"name": "operation_id"}
wait_for_operation_to_complete.return_value = None
with self.assertRaises(AirflowException) as cm:
self.gce_hook_no_project_id.stop_instance(zone=GCE_ZONE,
resource_id=GCE_INSTANCE)
stop_method.assert_not_called()
execute_method.assert_not_called()
err = cm.exception
self.assertIn("The project id must be passed", str(err))
wait_for_operation_to_complete.assert_not_called()
@mock.patch('airflow.gcp.hooks.compute.ComputeEngineHook.get_conn')
@mock.patch(
'airflow.gcp.hooks.compute.ComputeEngineHook._wait_for_operation_to_complete'
)
def test_set_machine_type_overridden_project_id(
self, wait_for_operation_to_complete, get_conn):
set_machine_type_method = get_conn.return_value.instances.return_value.setMachineType
execute_method = set_machine_type_method.return_value.execute
execute_method.return_value = {"name": "operation_id"}
wait_for_operation_to_complete.return_value = None
res = self.gce_hook_no_project_id.set_machine_type(
body={},
project_id='example-project',
zone=GCE_ZONE,
resource_id=GCE_INSTANCE)
self.assertIsNone(res)
set_machine_type_method.assert_called_once_with(
body={},
instance='instance',
project='example-project',
zone='zone')
execute_method.assert_called_once_with(num_retries=5)
wait_for_operation_to_complete.assert_called_once_with(
project_id='example-project',
operation_name='operation_id',
zone='zone')
@mock.patch('airflow.gcp.hooks.base.GoogleCloudBaseHook.project_id',
new_callable=PropertyMock,
return_value=None)
@mock.patch('airflow.gcp.hooks.compute.ComputeEngineHook.get_conn')
@mock.patch(
'airflow.gcp.hooks.compute.ComputeEngineHook._wait_for_operation_to_complete'
)
def test_set_machine_type_no_project_id(self,
wait_for_operation_to_complete,
get_conn, mock_project_id):
set_machine_type_method = get_conn.return_value.instances.return_value.setMachineType
execute_method = set_machine_type_method.return_value.execute
execute_method.return_value = {"name": "operation_id"}
wait_for_operation_to_complete.return_value = None
with self.assertRaises(AirflowException) as cm:
self.gce_hook_no_project_id.set_machine_type(
body={}, zone=GCE_ZONE, resource_id=GCE_INSTANCE)
set_machine_type_method.assert_not_called()
execute_method.assert_not_called()
err = cm.exception
self.assertIn("The project id must be passed", str(err))
wait_for_operation_to_complete.assert_not_called()
@mock.patch('airflow.gcp.hooks.compute.ComputeEngineHook.get_conn')
@mock.patch(
'airflow.gcp.hooks.compute.ComputeEngineHook._wait_for_operation_to_complete'
)
def test_get_instance_template_overridden_project_id(
self, wait_for_operation_to_complete, get_conn):
get_method = get_conn.return_value.instanceTemplates.return_value.get
execute_method = get_method.return_value.execute
execute_method.return_value = {"name": "operation_id"}
wait_for_operation_to_complete.return_value = None
res = self.gce_hook_no_project_id.get_instance_template(
resource_id=GCE_INSTANCE_TEMPLATE, project_id='example-project')
self.assertIsNotNone(res)
get_method.assert_called_once_with(
instanceTemplate='instance-template', project='example-project')
execute_method.assert_called_once_with(num_retries=5)
wait_for_operation_to_complete.assert_not_called()
@mock.patch('airflow.gcp.hooks.base.GoogleCloudBaseHook.project_id',
new_callable=PropertyMock,
return_value=None)
@mock.patch('airflow.gcp.hooks.compute.ComputeEngineHook.get_conn')
@mock.patch(
'airflow.gcp.hooks.compute.ComputeEngineHook._wait_for_operation_to_complete'
)
def test_get_instance_template_no_project_id(
self, wait_for_operation_to_complete, get_conn, mock_project_id):
get_method = get_conn.return_value.instanceTemplates.return_value.get
execute_method = get_method.return_value.execute
execute_method.return_value = {"name": "operation_id"}
wait_for_operation_to_complete.return_value = None
with self.assertRaises(AirflowException) as cm:
self.gce_hook_no_project_id.get_instance_template(
resource_id=GCE_INSTANCE_TEMPLATE)
get_method.assert_not_called()
execute_method.assert_not_called()
err = cm.exception
self.assertIn("The project id must be passed", str(err))
wait_for_operation_to_complete.assert_not_called()
@mock.patch('airflow.gcp.hooks.compute.ComputeEngineHook.get_conn')
@mock.patch(
'airflow.gcp.hooks.compute.ComputeEngineHook._wait_for_operation_to_complete'
)
def test_insert_instance_template_overridden_project_id(
self, wait_for_operation_to_complete, get_conn):
insert_method = get_conn.return_value.instanceTemplates.return_value.insert
execute_method = insert_method.return_value.execute
execute_method.return_value = {"name": "operation_id"}
wait_for_operation_to_complete.return_value = None
res = self.gce_hook_no_project_id.insert_instance_template(
project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST,
body={},
request_id=GCE_REQUEST_ID)
self.assertIsNone(res)
insert_method.assert_called_once_with(body={},
project='example-project',
requestId='request_id')
execute_method.assert_called_once_with(num_retries=5)
wait_for_operation_to_complete.assert_called_once_with(
project_id='example-project', operation_name='operation_id')
@mock.patch('airflow.gcp.hooks.base.GoogleCloudBaseHook.project_id',
new_callable=PropertyMock,
return_value=None)
@mock.patch('airflow.gcp.hooks.compute.ComputeEngineHook.get_conn')
@mock.patch(
'airflow.gcp.hooks.compute.ComputeEngineHook._wait_for_operation_to_complete'
)
def test_insert_instance_template_no_project_id(
self, wait_for_operation_to_complete, get_conn, mock_project_id):
insert_method = get_conn.return_value.instanceTemplates.return_value.insert
execute_method = insert_method.return_value.execute
execute_method.return_value = {"name": "operation_id"}
wait_for_operation_to_complete.return_value = None
with self.assertRaises(AirflowException) as cm:
self.gce_hook_no_project_id.insert_instance_template(
body={}, request_id=GCE_REQUEST_ID)
insert_method.assert_not_called()
execute_method.assert_not_called()
err = cm.exception
self.assertIn("The project id must be passed", str(err))
wait_for_operation_to_complete.assert_not_called()
@mock.patch('airflow.gcp.hooks.compute.ComputeEngineHook.get_conn')
@mock.patch(
'airflow.gcp.hooks.compute.ComputeEngineHook._wait_for_operation_to_complete'
)
def test_get_instance_group_manager_overridden_project_id(
self, wait_for_operation_to_complete, get_conn):
get_method = get_conn.return_value.instanceGroupManagers.return_value.get
execute_method = get_method.return_value.execute
execute_method.return_value = {"name": "operation_id"}
wait_for_operation_to_complete.return_value = None
res = self.gce_hook_no_project_id.get_instance_group_manager(
project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST,
zone=GCE_ZONE,
resource_id=GCE_INSTANCE_GROUP_MANAGER)
self.assertIsNotNone(res)
get_method.assert_called_once_with(
instanceGroupManager='instance_group_manager',
project='example-project',
zone='zone')
execute_method.assert_called_once_with(num_retries=5)
wait_for_operation_to_complete.assert_not_called()
@mock.patch('airflow.gcp.hooks.base.GoogleCloudBaseHook.project_id',
new_callable=PropertyMock,
return_value=None)
@mock.patch('airflow.gcp.hooks.compute.ComputeEngineHook.get_conn')
@mock.patch(
'airflow.gcp.hooks.compute.ComputeEngineHook._wait_for_operation_to_complete'
)
def test_get_instance_group_manager_no_project_id(
self, wait_for_operation_to_complete, get_conn, mock_project_id):
get_method = get_conn.return_value.instanceGroupManagers.return_value.get
execute_method = get_method.return_value.execute
execute_method.return_value = {"name": "operation_id"}
wait_for_operation_to_complete.return_value = None
with self.assertRaises(AirflowException) as cm:
self.gce_hook_no_project_id.get_instance_group_manager(
zone=GCE_ZONE, resource_id=GCE_INSTANCE_GROUP_MANAGER)
get_method.assert_not_called()
execute_method.assert_not_called()
err = cm.exception
self.assertIn("The project id must be passed", str(err))
wait_for_operation_to_complete.assert_not_called()
@mock.patch('airflow.gcp.hooks.compute.ComputeEngineHook.get_conn')
@mock.patch(
'airflow.gcp.hooks.compute.ComputeEngineHook._wait_for_operation_to_complete'
)
def test_patch_instance_group_manager_overridden_project_id(
self, wait_for_operation_to_complete, get_conn):
patch_method = get_conn.return_value.instanceGroupManagers.return_value.patch
execute_method = patch_method.return_value.execute
execute_method.return_value = {"name": "operation_id"}
wait_for_operation_to_complete.return_value = None
res = self.gce_hook_no_project_id.patch_instance_group_manager(
project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST,
zone=GCE_ZONE,
resource_id=GCE_INSTANCE_GROUP_MANAGER,
body={},
request_id=GCE_REQUEST_ID)
self.assertIsNone(res)
patch_method.assert_called_once_with(
body={},
instanceGroupManager='instance_group_manager',
project='example-project',
requestId='request_id',
zone='zone')
execute_method.assert_called_once_with(num_retries=5)
wait_for_operation_to_complete.assert_called_once_with(
operation_name='operation_id',
project_id='example-project',
zone='zone')
@mock.patch('airflow.gcp.hooks.base.GoogleCloudBaseHook.project_id',
new_callable=PropertyMock,
return_value=None)
@mock.patch('airflow.gcp.hooks.compute.ComputeEngineHook.get_conn')
@mock.patch(
'airflow.gcp.hooks.compute.ComputeEngineHook._wait_for_operation_to_complete'
)
def test_patch_instance_group_manager_no_project_id(
self, wait_for_operation_to_complete, get_conn, mock_project_id):
patch_method = get_conn.return_value.instanceGroupManagers.return_value.patch
execute_method = patch_method.return_value.execute
execute_method.return_value = {"name": "operation_id"}
wait_for_operation_to_complete.return_value = None
with self.assertRaises(AirflowException) as cm:
self.gce_hook_no_project_id.patch_instance_group_manager(
zone=GCE_ZONE,
resource_id=GCE_INSTANCE_GROUP_MANAGER,
body={},
request_id=GCE_REQUEST_ID)
patch_method.assert_not_called()
execute_method.assert_not_called()
err = cm.exception
self.assertIn("The project id must be passed", str(err))
wait_for_operation_to_complete.assert_not_called()
class TestGcpComputeHookDefaultProjectId(unittest.TestCase):
def setUp(self):
with mock.patch('airflow.gcp.hooks.base.GoogleCloudBaseHook.__init__',
new=mock_base_gcp_hook_default_project_id):
self.gce_hook = ComputeEngineHook(gcp_conn_id='test')
@mock.patch('airflow.gcp.hooks.base.GoogleCloudBaseHook.project_id',
new_callable=PropertyMock,
return_value=GCP_PROJECT_ID_HOOK_UNIT_TEST)
@mock.patch('airflow.gcp.hooks.compute.ComputeEngineHook.get_conn')
@mock.patch(
'airflow.gcp.hooks.compute.ComputeEngineHook._wait_for_operation_to_complete'
)
def test_start_instance(self, wait_for_operation_to_complete, get_conn,
mock_project_id):
start_method = get_conn.return_value.instances.return_value.start
execute_method = start_method.return_value.execute
execute_method.return_value = {"name": "operation_id"}
wait_for_operation_to_complete.return_value = None
res = self.gce_hook.start_instance(zone=GCE_ZONE,
resource_id=GCE_INSTANCE)
self.assertIsNone(res)
start_method.assert_called_once_with(instance='instance',
project='example-project',
zone='zone')
execute_method.assert_called_once_with(num_retries=5)
wait_for_operation_to_complete.assert_called_once_with(
project_id='example-project',
operation_name='operation_id',
zone='zone')
@mock.patch('airflow.gcp.hooks.compute.ComputeEngineHook.get_conn')
@mock.patch(
'airflow.gcp.hooks.compute.ComputeEngineHook._wait_for_operation_to_complete'
)
def test_start_instance_overridden_project_id(
self, wait_for_operation_to_complete, get_conn):
start_method = get_conn.return_value.instances.return_value.start
execute_method = start_method.return_value.execute
execute_method.return_value = {"name": "operation_id"}
wait_for_operation_to_complete.return_value = None
res = self.gce_hook.start_instance(project_id='new-project',
zone=GCE_ZONE,
resource_id=GCE_INSTANCE)
self.assertIsNone(res)
start_method.assert_called_once_with(instance='instance',
project='new-project',
zone='zone')
execute_method.assert_called_once_with(num_retries=5)
wait_for_operation_to_complete.assert_called_once_with(
project_id='new-project',
operation_name='operation_id',
zone='zone')
@mock.patch('airflow.gcp.hooks.base.GoogleCloudBaseHook.project_id',
new_callable=PropertyMock,
return_value=GCP_PROJECT_ID_HOOK_UNIT_TEST)
@mock.patch('airflow.gcp.hooks.compute.ComputeEngineHook.get_conn')
@mock.patch(
'airflow.gcp.hooks.compute.ComputeEngineHook._wait_for_operation_to_complete'
)
def test_stop_instance(self, wait_for_operation_to_complete, get_conn,
mock_project_id):
stop_method = get_conn.return_value.instances.return_value.stop
execute_method = stop_method.return_value.execute
execute_method.return_value = {"name": "operation_id"}
wait_for_operation_to_complete.return_value = None
res = self.gce_hook.stop_instance(zone=GCE_ZONE,
resource_id=GCE_INSTANCE)
self.assertIsNone(res)
stop_method.assert_called_once_with(instance='instance',
project='example-project',
zone='zone')
execute_method.assert_called_once_with(num_retries=5)
wait_for_operation_to_complete.assert_called_once_with(
project_id='example-project',
operation_name='operation_id',
zone='zone')
@mock.patch('airflow.gcp.hooks.compute.ComputeEngineHook.get_conn')
@mock.patch(
'airflow.gcp.hooks.compute.ComputeEngineHook._wait_for_operation_to_complete'
)
def test_stop_instance_overridden_project_id(
self, wait_for_operation_to_complete, get_conn):
stop_method = get_conn.return_value.instances.return_value.stop
execute_method = stop_method.return_value.execute
execute_method.return_value = {"name": "operation_id"}
wait_for_operation_to_complete.return_value = None
res = self.gce_hook.stop_instance(project_id='new-project',
zone=GCE_ZONE,
resource_id=GCE_INSTANCE)
self.assertIsNone(res)
stop_method.assert_called_once_with(instance='instance',
project='new-project',
zone='zone')
execute_method.assert_called_once_with(num_retries=5)
wait_for_operation_to_complete.assert_called_once_with(
project_id='new-project',
operation_name='operation_id',
zone='zone')
@mock.patch('airflow.gcp.hooks.base.GoogleCloudBaseHook.project_id',
new_callable=PropertyMock,
return_value=GCP_PROJECT_ID_HOOK_UNIT_TEST)
@mock.patch('airflow.gcp.hooks.compute.ComputeEngineHook.get_conn')
@mock.patch(
'airflow.gcp.hooks.compute.ComputeEngineHook._wait_for_operation_to_complete'
)
def test_set_machine_type_instance(self, wait_for_operation_to_complete,
get_conn, mock_project_id):
execute_method = get_conn.return_value.instances.return_value.setMachineType.return_value.execute
execute_method.return_value = {"name": "operation_id"}
wait_for_operation_to_complete.return_value = None
res = self.gce_hook.set_machine_type(body={},
zone=GCE_ZONE,
resource_id=GCE_INSTANCE)
self.assertIsNone(res)
execute_method.assert_called_once_with(num_retries=5)
wait_for_operation_to_complete.assert_called_once_with(
project_id='example-project',
operation_name='operation_id',
zone='zone')
@mock.patch('airflow.gcp.hooks.compute.ComputeEngineHook.get_conn')
@mock.patch(
'airflow.gcp.hooks.compute.ComputeEngineHook._wait_for_operation_to_complete'
)
def test_set_machine_type_instance_overridden_project_id(
self, wait_for_operation_to_complete, get_conn):
execute_method = get_conn.return_value.instances.return_value.setMachineType.return_value.execute
execute_method.return_value = {"name": "operation_id"}
wait_for_operation_to_complete.return_value = None
res = self.gce_hook.set_machine_type(project_id='new-project',
body={},
zone=GCE_ZONE,
resource_id=GCE_INSTANCE)
self.assertIsNone(res)
execute_method.assert_called_once_with(num_retries=5)
wait_for_operation_to_complete.assert_called_once_with(
project_id='new-project',
operation_name='operation_id',
zone='zone')
@mock.patch('airflow.gcp.hooks.base.GoogleCloudBaseHook.project_id',
new_callable=PropertyMock,
return_value=GCP_PROJECT_ID_HOOK_UNIT_TEST)
@mock.patch('airflow.gcp.hooks.compute.ComputeEngineHook.get_conn')
@mock.patch(
'airflow.gcp.hooks.compute.ComputeEngineHook._wait_for_operation_to_complete'
)
def test_get_instance_template(self, wait_for_operation_to_complete,
get_conn, mock_project_id):
get_method = get_conn.return_value.instanceTemplates.return_value.get
execute_method = get_method.return_value.execute
execute_method.return_value = {"name": "operation_id"}
wait_for_operation_to_complete.return_value = None
res = self.gce_hook.get_instance_template(
resource_id=GCE_INSTANCE_TEMPLATE)
self.assertIsNotNone(res)
get_method.assert_called_once_with(
instanceTemplate='instance-template', project='example-project')
execute_method.assert_called_once_with(num_retries=5)
wait_for_operation_to_complete.assert_not_called()
@mock.patch('airflow.gcp.hooks.compute.ComputeEngineHook.get_conn')
@mock.patch(
'airflow.gcp.hooks.compute.ComputeEngineHook._wait_for_operation_to_complete'
)
def test_get_instance_template_overridden_project_id(
self, wait_for_operation_to_complete, get_conn):
get_method = get_conn.return_value.instanceTemplates.return_value.get
execute_method = get_method.return_value.execute
execute_method.return_value = {"name": "operation_id"}
wait_for_operation_to_complete.return_value = None
res = self.gce_hook.get_instance_template(
project_id='new-project', resource_id=GCE_INSTANCE_TEMPLATE)
self.assertIsNotNone(res)
get_method.assert_called_once_with(
instanceTemplate='instance-template', project='new-project')
execute_method.assert_called_once_with(num_retries=5)
wait_for_operation_to_complete.assert_not_called()
@mock.patch('airflow.gcp.hooks.base.GoogleCloudBaseHook.project_id',
new_callable=PropertyMock,
return_value=GCP_PROJECT_ID_HOOK_UNIT_TEST)
@mock.patch('airflow.gcp.hooks.compute.ComputeEngineHook.get_conn')
@mock.patch(
'airflow.gcp.hooks.compute.ComputeEngineHook._wait_for_operation_to_complete'
)
def test_insert_instance_template(self, wait_for_operation_to_complete,
get_conn, mock_project_id):
insert_method = get_conn.return_value.instanceTemplates.return_value.insert
execute_method = insert_method.return_value.execute
execute_method.return_value = {"name": "operation_id"}
wait_for_operation_to_complete.return_value = None
res = self.gce_hook.insert_instance_template(body={},
request_id=GCE_REQUEST_ID)
self.assertIsNone(res)
insert_method.assert_called_once_with(body={},
project='example-project',
requestId='request_id')
execute_method.assert_called_once_with(num_retries=5)
wait_for_operation_to_complete.assert_called_once_with(
project_id='example-project', operation_name='operation_id')
@mock.patch('airflow.gcp.hooks.compute.ComputeEngineHook.get_conn')
@mock.patch(
'airflow.gcp.hooks.compute.ComputeEngineHook._wait_for_operation_to_complete'
)
def test_insert_instance_template_overridden_project_id(
self, wait_for_operation_to_complete, get_conn):
insert_method = get_conn.return_value.instanceTemplates.return_value.insert
execute_method = insert_method.return_value.execute
execute_method.return_value = {"name": "operation_id"}
wait_for_operation_to_complete.return_value = None
res = self.gce_hook.insert_instance_template(project_id='new-project',
body={},
request_id=GCE_REQUEST_ID)
self.assertIsNone(res)
insert_method.assert_called_once_with(body={},
project='new-project',
requestId='request_id')
execute_method.assert_called_once_with(num_retries=5)
wait_for_operation_to_complete.assert_called_once_with(
project_id='new-project', operation_name='operation_id')
@mock.patch('airflow.gcp.hooks.base.GoogleCloudBaseHook.project_id',
new_callable=PropertyMock,
return_value=GCP_PROJECT_ID_HOOK_UNIT_TEST)
@mock.patch('airflow.gcp.hooks.compute.ComputeEngineHook.get_conn')
@mock.patch(
'airflow.gcp.hooks.compute.ComputeEngineHook._wait_for_operation_to_complete'
)
def test_get_instance_group_manager(self, wait_for_operation_to_complete,
get_conn, mock_project_id):
get_method = get_conn.return_value.instanceGroupManagers.return_value.get
execute_method = get_method.return_value.execute
execute_method.return_value = {"name": "operation_id"}
wait_for_operation_to_complete.return_value = None
res = self.gce_hook.get_instance_group_manager(
zone=GCE_ZONE, resource_id=GCE_INSTANCE_GROUP_MANAGER)
self.assertIsNotNone(res)
get_method.assert_called_once_with(
instanceGroupManager='instance_group_manager',
project='example-project',
zone='zone')
execute_method.assert_called_once_with(num_retries=5)
wait_for_operation_to_complete.assert_not_called()
@mock.patch('airflow.gcp.hooks.compute.ComputeEngineHook.get_conn')
@mock.patch(
'airflow.gcp.hooks.compute.ComputeEngineHook._wait_for_operation_to_complete'
)
def test_get_instance_group_manager_overridden_project_id(
self, wait_for_operation_to_complete, get_conn):
get_method = get_conn.return_value.instanceGroupManagers.return_value.get
execute_method = get_method.return_value.execute
execute_method.return_value = {"name": "operation_id"}
wait_for_operation_to_complete.return_value = None
res = self.gce_hook.get_instance_group_manager(
project_id='new-project',
zone=GCE_ZONE,
resource_id=GCE_INSTANCE_GROUP_MANAGER)
self.assertIsNotNone(res)
get_method.assert_called_once_with(
instanceGroupManager='instance_group_manager',
project='new-project',
zone='zone')
execute_method.assert_called_once_with(num_retries=5)
wait_for_operation_to_complete.assert_not_called()
@mock.patch('airflow.gcp.hooks.base.GoogleCloudBaseHook.project_id',
new_callable=PropertyMock,
return_value=GCP_PROJECT_ID_HOOK_UNIT_TEST)
@mock.patch('airflow.gcp.hooks.compute.ComputeEngineHook.get_conn')
@mock.patch(
'airflow.gcp.hooks.compute.ComputeEngineHook._wait_for_operation_to_complete'
)
def test_patch_instance_group_manager(self, wait_for_operation_to_complete,
get_conn, mock_project_id):
patch_method = get_conn.return_value.instanceGroupManagers.return_value.patch
execute_method = patch_method.return_value.execute
execute_method.return_value = {"name": "operation_id"}
wait_for_operation_to_complete.return_value = None
res = self.gce_hook.patch_instance_group_manager(
zone=GCE_ZONE,
resource_id=GCE_INSTANCE_GROUP_MANAGER,
body={},
request_id=GCE_REQUEST_ID)
self.assertIsNone(res)
patch_method.assert_called_once_with(
body={},
instanceGroupManager='instance_group_manager',
project='example-project',
requestId='request_id',
zone='zone')
execute_method.assert_called_once_with(num_retries=5)
wait_for_operation_to_complete.assert_called_once_with(
operation_name='operation_id',
project_id='example-project',
zone='zone')
@mock.patch('airflow.gcp.hooks.compute.ComputeEngineHook.get_conn')
@mock.patch(
'airflow.gcp.hooks.compute.ComputeEngineHook._wait_for_operation_to_complete'
)
def test_patch_instance_group_manager_overridden_project_id(
self, wait_for_operation_to_complete, get_conn):
patch_method = get_conn.return_value.instanceGroupManagers.return_value.patch
execute_method = patch_method.return_value.execute
execute_method.return_value = {"name": "operation_id"}
wait_for_operation_to_complete.return_value = None
res = self.gce_hook.patch_instance_group_manager(
project_id='new-project',
zone=GCE_ZONE,
resource_id=GCE_INSTANCE_GROUP_MANAGER,
body={},
request_id=GCE_REQUEST_ID)
self.assertIsNone(res)
patch_method.assert_called_once_with(
body={},
instanceGroupManager='instance_group_manager',
project='new-project',
requestId='request_id',
zone='zone')
execute_method.assert_called_once_with(num_retries=5)
wait_for_operation_to_complete.assert_called_once_with(
operation_name='operation_id',
project_id='new-project',
zone='zone')
@mock.patch('airflow.gcp.hooks.compute.ComputeEngineHook.get_conn')
@mock.patch(
'airflow.gcp.hooks.compute.ComputeEngineHook._check_global_operation_status'
)
def test_wait_for_operation_to_complete_no_zone(self,
mock_operation_status,
mock_get_conn):
service = "test-service"
project_id = "test-project"
operation_name = "test-operation"
num_retries = self.gce_hook.num_retries
# Test success
mock_get_conn.return_value = service
mock_operation_status.return_value = {
'status': GceOperationStatus.DONE,
'error': None
}
self.gce_hook._wait_for_operation_to_complete(
project_id=project_id, operation_name=operation_name, zone=None)
mock_operation_status.assert_called_once_with(
service=service,
operation_name=operation_name,
project_id=project_id,
num_retries=num_retries)
@mock.patch('airflow.gcp.hooks.compute.ComputeEngineHook.get_conn')
@mock.patch(
'airflow.gcp.hooks.compute.ComputeEngineHook._check_global_operation_status'
)
def test_wait_for_operation_to_complete_no_zone_error(
self, mock_operation_status, mock_get_conn):
service = "test-service"
project_id = "test-project"
operation_name = "test-operation"
# Test error
mock_get_conn.return_value = service
mock_operation_status.return_value = {
'status': GceOperationStatus.DONE,
'error': {
'errors': "some nasty errors"
},
'httpErrorStatusCode': 400,
'httpErrorMessage': 'sample msg'
}
with self.assertRaises(AirflowException):
self.gce_hook._wait_for_operation_to_complete(
project_id=project_id,
operation_name=operation_name,
zone=None)
@mock.patch('airflow.gcp.hooks.compute.ComputeEngineHook.get_conn')
@mock.patch(
'airflow.gcp.hooks.compute.ComputeEngineHook._check_zone_operation_status'
)
def test_wait_for_operation_to_complete_with_zone(self,
mock_operation_status,
mock_get_conn):
service = "test-service"
project_id = "test-project"
operation_name = "test-operation"
zone = 'west-europe3'
num_retries = self.gce_hook.num_retries
# Test success
mock_get_conn.return_value = service
mock_operation_status.return_value = {
'status': GceOperationStatus.DONE,
'error': None
}
self.gce_hook._wait_for_operation_to_complete(
project_id=project_id, operation_name=operation_name, zone=zone)
mock_operation_status.assert_called_once_with(service, operation_name,
project_id, zone,
num_retries)