def test_different_error_reraised(self, mock_hook):
op = BigtableUpdateInstanceOperator(
project_id=PROJECT_ID,
instance_id=INSTANCE_ID,
instance_display_name=INSTANCE_DISPLAY_NAME,
instance_type=INSTANCE_TYPE,
instance_labels=INSTANCE_LABELS,
task_id="id",
gcp_conn_id=GCP_CONN_ID,
impersonation_chain=IMPERSONATION_CHAIN,
)
mock_hook.return_value.update_instance.side_effect = mock.Mock(
side_effect=google.api_core.exceptions.GoogleAPICallError('error'))
with self.assertRaises(google.api_core.exceptions.GoogleAPICallError):
op.execute(None)
mock_hook.assert_called_once_with(
gcp_conn_id=GCP_CONN_ID,
impersonation_chain=IMPERSONATION_CHAIN,
)
mock_hook.return_value.update_instance.assert_called_once_with(
project_id=PROJECT_ID,
instance_id=INSTANCE_ID,
instance_display_name=INSTANCE_DISPLAY_NAME,
instance_type=INSTANCE_TYPE,
instance_labels=INSTANCE_LABELS,
timeout=None,
)
def test_update_instance_that_doesnt_exists(self, mock_hook):
mock_hook.return_value.get_instance.return_value = None
with pytest.raises(AirflowException) as ctx:
op = BigtableUpdateInstanceOperator(
project_id=PROJECT_ID,
instance_id=INSTANCE_ID,
instance_display_name=INSTANCE_DISPLAY_NAME,
instance_type=INSTANCE_TYPE,
instance_labels=INSTANCE_LABELS,
task_id="id",
gcp_conn_id=GCP_CONN_ID,
impersonation_chain=IMPERSONATION_CHAIN,
)
op.execute(None)
err = ctx.value
assert str(
err) == f"Dependency: instance '{INSTANCE_ID}' does not exist."
mock_hook.assert_called_once_with(
gcp_conn_id=GCP_CONN_ID,
impersonation_chain=IMPERSONATION_CHAIN,
)
mock_hook.return_value.update_instance.assert_not_called()
def test_update_instance_that_doesnt_exists_empty_project_id(
self, mock_hook):
mock_hook.return_value.get_instance.return_value = None
with self.assertRaises(AirflowException) as e:
op = BigtableUpdateInstanceOperator(
instance_id=INSTANCE_ID,
instance_display_name=INSTANCE_DISPLAY_NAME,
instance_type=INSTANCE_TYPE,
instance_labels=INSTANCE_LABELS,
task_id="id",
gcp_conn_id=GCP_CONN_ID,
impersonation_chain=IMPERSONATION_CHAIN,
)
op.execute(None)
err = e.exception
self.assertEqual(
str(err),
"Dependency: instance '{}' does not exist.".format(INSTANCE_ID))
mock_hook.assert_called_once_with(
gcp_conn_id=GCP_CONN_ID,
impersonation_chain=IMPERSONATION_CHAIN,
)
mock_hook.return_value.update_instance.assert_not_called()
def test_empty_attribute(self, missing_attribute, project_id, instance_id, mock_hook):
with self.assertRaises(AirflowException) as e:
BigtableUpdateInstanceOperator(
project_id=project_id,
instance_id=instance_id,
instance_display_name=INSTANCE_DISPLAY_NAME,
instance_type=INSTANCE_TYPE,
instance_labels=INSTANCE_LABELS,
task_id="id",
)
err = e.exception
self.assertEqual(str(err), 'Empty parameter: {}'.format(missing_attribute))
mock_hook.assert_not_called()
def test_empty_attribute(self, missing_attribute, project_id, instance_id,
mock_hook):
with pytest.raises(AirflowException) as ctx:
BigtableUpdateInstanceOperator(
project_id=project_id,
instance_id=instance_id,
instance_display_name=INSTANCE_DISPLAY_NAME,
instance_type=INSTANCE_TYPE,
instance_labels=INSTANCE_LABELS,
task_id="id",
)
err = ctx.value
assert str(err) == f'Empty parameter: {missing_attribute}'
mock_hook.assert_not_called()
def test_update_execute_empty_project_id(self, mock_hook):
op = BigtableUpdateInstanceOperator(
instance_id=INSTANCE_ID,
instance_display_name=INSTANCE_DISPLAY_NAME,
instance_type=INSTANCE_TYPE,
instance_labels=INSTANCE_LABELS,
task_id="id",
gcp_conn_id=GCP_CONN_ID)
op.execute(None)
mock_hook.assert_called_once_with(gcp_conn_id=GCP_CONN_ID)
mock_hook.return_value.update_instance.assert_called_once_with(
project_id=None,
instance_id=INSTANCE_ID,
instance_display_name=INSTANCE_DISPLAY_NAME,
instance_type=INSTANCE_TYPE,
instance_labels=INSTANCE_LABELS,
timeout=None)
def test_delete_execute(self, mock_hook):
op = BigtableUpdateInstanceOperator(
project_id=PROJECT_ID,
instance_id=INSTANCE_ID,
instance_display_name=INSTANCE_DISPLAY_NAME,
instance_type=INSTANCE_TYPE,
instance_labels=INSTANCE_LABELS,
task_id="id",
gcp_conn_id=GCP_CONN_ID,
impersonation_chain=IMPERSONATION_CHAIN,
)
op.execute(None)
mock_hook.assert_called_once_with(
gcp_conn_id=GCP_CONN_ID,
impersonation_chain=IMPERSONATION_CHAIN,
)
mock_hook.return_value.update_instance.assert_called_once_with(
project_id=PROJECT_ID,
instance_id=INSTANCE_ID,
instance_display_name=INSTANCE_DISPLAY_NAME,
instance_type=INSTANCE_TYPE,
instance_labels=INSTANCE_LABELS,
timeout=None)
main_cluster_zone=CBT_CLUSTER_ZONE,
instance_display_name=CBT_INSTANCE_DISPLAY_NAME,
instance_type=int(CBT_INSTANCE_TYPE),
instance_labels=json.loads(CBT_INSTANCE_LABELS),
cluster_nodes=int(CBT_CLUSTER_NODES),
cluster_storage_type=int(CBT_CLUSTER_STORAGE_TYPE),
task_id='create_instance_task2',
)
create_instance_task >> create_instance_task2
# [END howto_operator_gcp_bigtable_instance_create]
# [START howto_operator_gcp_bigtable_instance_update]
update_instance_task = BigtableUpdateInstanceOperator(
instance_id=CBT_INSTANCE_ID,
instance_display_name=CBT_INSTANCE_DISPLAY_NAME_UPDATED,
instance_type=int(CBT_INSTANCE_TYPE_PROD),
instance_labels=json.loads(CBT_INSTANCE_LABELS_UPDATED),
task_id='update_instance_task',
)
# [END howto_operator_gcp_bigtable_instance_update]
# [START howto_operator_gcp_bigtable_cluster_update]
cluster_update_task = BigtableUpdateClusterOperator(
project_id=GCP_PROJECT_ID,
instance_id=CBT_INSTANCE_ID,
cluster_id=CBT_CLUSTER_ID,
nodes=int(CBT_CLUSTER_NODES_UPDATED),
task_id='update_cluster_task',
)
cluster_update_task2 = BigtableUpdateClusterOperator(
instance_id=CBT_INSTANCE_ID,