Ejemplo n.º 1
0
    def test_different_error_reraised(self, mock_hook):
        op = BigtableUpdateClusterOperator(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)
Ejemplo n.º 2
0
    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 = BigtableUpdateClusterOperator(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()
Ejemplo n.º 3
0
 def test_empty_attribute(self, missing_attribute, project_id, instance_id,
                          cluster_id, nodes, mock_hook):
     with self.assertRaises(AirflowException) as e:
         BigtableUpdateClusterOperator(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()
Ejemplo n.º 4
0
    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 = BigtableUpdateClusterOperator(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)
Ejemplo n.º 5
0
        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 = 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,
        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 = BigtableDeleteInstanceOperator(
        project_id=GCP_PROJECT_ID,