def test_different_error_reraised(self, mock_hook): op = BigtableClusterUpdateOperator(project_id=PROJECT_ID, instance_id=INSTANCE_ID, cluster_id=CLUSTER_ID, nodes=NODES, task_id="id", gcp_conn_id=GCP_CONN_ID) instance = mock_hook.return_value.get_instance.return_value = mock.Mock( Instance) mock_hook.return_value.update_cluster.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) mock_hook.return_value.update_cluster.assert_called_once_with( instance=instance, cluster_id=CLUSTER_ID, nodes=NODES)
def test_updating_cluster_but_instance_does_not_exists_empty_project_id( self, mock_hook): mock_hook.return_value.get_instance.return_value = None with self.assertRaises(AirflowException) as e: op = BigtableClusterUpdateOperator(instance_id=INSTANCE_ID, cluster_id=CLUSTER_ID, nodes=NODES, task_id="id", gcp_conn_id=GCP_CONN_ID) 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) mock_hook.return_value.update_cluster.assert_not_called()
def test_empty_attribute(self, missing_attribute, project_id, instance_id, cluster_id, nodes, mock_hook): with self.assertRaises(AirflowException) as e: BigtableClusterUpdateOperator(project_id=project_id, instance_id=instance_id, cluster_id=cluster_id, nodes=nodes, task_id="id", gcp_conn_id=GCP_CONN_ID) err = e.exception self.assertEqual(str(err), 'Empty parameter: {}'.format(missing_attribute)) mock_hook.assert_not_called()
def test_updating_cluster_that_does_not_exists_empty_project_id( self, mock_hook): instance = mock_hook.return_value.get_instance.return_value = mock.Mock( Instance) mock_hook.return_value.update_cluster.side_effect = mock.Mock( side_effect=google.api_core.exceptions.NotFound( "Cluster not found.")) with self.assertRaises(AirflowException) as e: op = BigtableClusterUpdateOperator(instance_id=INSTANCE_ID, cluster_id=CLUSTER_ID, nodes=NODES, task_id="id", gcp_conn_id=GCP_CONN_ID) op.execute(None) err = e.exception self.assertEqual( str(err), "Dependency: cluster '{}' does not exist for instance '{}'.". format(CLUSTER_ID, INSTANCE_ID)) mock_hook.assert_called_once_with(gcp_conn_id=GCP_CONN_ID) mock_hook.return_value.update_cluster.assert_called_once_with( instance=instance, cluster_id=CLUSTER_ID, nodes=NODES)
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_cluster_update] cluster_update_task = BigtableClusterUpdateOperator( 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 = BigtableClusterUpdateOperator( instance_id=CBT_INSTANCE_ID, cluster_id=CBT_CLUSTER_ID, nodes=int(CBT_CLUSTER_NODES_UPDATED), task_id='update_cluster_task2', ) cluster_update_task >> cluster_update_task2 # [END howto_operator_gcp_bigtable_cluster_update] # [START howto_operator_gcp_bigtable_instance_delete] delete_instance_task = BigtableInstanceDeleteOperator( project_id=GCP_PROJECT_ID,