Exemplo n.º 1
0
 def test_delete_execute(self, mock_hook):
     op = BigtableDeleteTableOperator(project_id=PROJECT_ID,
                                      instance_id=INSTANCE_ID,
                                      table_id=TABLE_ID,
                                      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.delete_table.assert_called_once_with(
         project_id=PROJECT_ID, instance_id=INSTANCE_ID, table_id=TABLE_ID)
Exemplo n.º 2
0
    def test_deleting_table_that_doesnt_exists_empty_project_id(
            self, mock_hook):
        op = BigtableDeleteTableOperator(instance_id=INSTANCE_ID,
                                         table_id=TABLE_ID,
                                         task_id="id",
                                         gcp_conn_id=GCP_CONN_ID)

        mock_hook.return_value.delete_table.side_effect = mock.Mock(
            side_effect=google.api_core.exceptions.NotFound(
                "Table not found."))
        op.execute(None)
        mock_hook.assert_called_once_with(gcp_conn_id=GCP_CONN_ID)
        mock_hook.return_value.delete_table.assert_called_once_with(
            project_id=None, instance_id=INSTANCE_ID, table_id=TABLE_ID)
Exemplo n.º 3
0
    def test_different_error_reraised(self, mock_hook):
        op = BigtableDeleteTableOperator(project_id=PROJECT_ID,
                                         instance_id=INSTANCE_ID,
                                         table_id=TABLE_ID,
                                         task_id="id",
                                         gcp_conn_id=GCP_CONN_ID)
        mock_hook.return_value.delete_table.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.delete_table.assert_called_once_with(
            project_id=PROJECT_ID, instance_id=INSTANCE_ID, table_id=TABLE_ID)
Exemplo n.º 4
0
    def test_deleting_table_when_instance_doesnt_exists(self, mock_hook):
        op = BigtableDeleteTableOperator(project_id=PROJECT_ID,
                                         instance_id=INSTANCE_ID,
                                         table_id=TABLE_ID,
                                         task_id="id",
                                         gcp_conn_id=GCP_CONN_ID)

        mock_hook.return_value.get_instance.return_value = None
        with self.assertRaises(AirflowException) as e:
            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.delete_table.assert_not_called()
Exemplo n.º 5
0
    def test_deleting_table_that_doesnt_exists(self, mock_hook):
        op = BigtableDeleteTableOperator(
            project_id=PROJECT_ID,
            instance_id=INSTANCE_ID,
            table_id=TABLE_ID,
            task_id="id",
            gcp_conn_id=GCP_CONN_ID,
            impersonation_chain=IMPERSONATION_CHAIN,
        )

        mock_hook.return_value.delete_table.side_effect = mock.Mock(
            side_effect=google.api_core.exceptions.NotFound(
                "Table not found."))
        op.execute(None)
        mock_hook.assert_called_once_with(
            gcp_conn_id=GCP_CONN_ID,
            impersonation_chain=IMPERSONATION_CHAIN,
        )
        mock_hook.return_value.delete_table.assert_called_once_with(
            project_id=PROJECT_ID, instance_id=INSTANCE_ID, table_id=TABLE_ID)
Exemplo n.º 6
0
 def test_empty_attribute(self, missing_attribute, project_id, instance_id,
                          table_id, mock_hook):
     with self.assertRaises(AirflowException) as e:
         BigtableDeleteTableOperator(project_id=project_id,
                                     instance_id=instance_id,
                                     table_id=table_id,
                                     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()
Exemplo n.º 7
0
    def test_deleting_table_when_instance_doesnt_exists(self, mock_hook):
        op = BigtableDeleteTableOperator(
            project_id=PROJECT_ID,
            instance_id=INSTANCE_ID,
            table_id=TABLE_ID,
            task_id="id",
            gcp_conn_id=GCP_CONN_ID,
            impersonation_chain=IMPERSONATION_CHAIN,
        )

        mock_hook.return_value.get_instance.return_value = None
        with pytest.raises(AirflowException) as ctx:
            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.delete_table.assert_not_called()
Exemplo n.º 8
0
 def test_empty_attribute(self, missing_attribute, project_id, instance_id,
                          table_id, mock_hook):
     with pytest.raises(AirflowException) as ctx:
         BigtableDeleteTableOperator(
             project_id=project_id,
             instance_id=instance_id,
             table_id=table_id,
             task_id="id",
             gcp_conn_id=GCP_CONN_ID,
         )
     err = ctx.value
     assert str(err) == f'Empty parameter: {missing_attribute}'
     mock_hook.assert_not_called()
Exemplo n.º 9
0
        timeout=180,
        task_id='wait_for_table_replication_task',
    )
    wait_for_table_replication_task2 = BigtableTableReplicationCompletedSensor(
        instance_id=CBT_INSTANCE_ID,
        table_id=CBT_TABLE_ID,
        poke_interval=int(CBT_POKE_INTERVAL),
        timeout=180,
        task_id='wait_for_table_replication_task2',
    )
    # [END howto_operator_gcp_bigtable_table_wait_for_replication]

    # [START howto_operator_gcp_bigtable_table_delete]
    delete_table_task = BigtableDeleteTableOperator(
        project_id=GCP_PROJECT_ID,
        instance_id=CBT_INSTANCE_ID,
        table_id=CBT_TABLE_ID,
        task_id='delete_table_task',
    )
    delete_table_task2 = BigtableDeleteTableOperator(
        instance_id=CBT_INSTANCE_ID,
        table_id=CBT_TABLE_ID,
        task_id='delete_table_task2',
    )
    # [END howto_operator_gcp_bigtable_table_delete]

    wait_for_table_replication_task >> delete_table_task
    wait_for_table_replication_task2 >> delete_table_task
    wait_for_table_replication_task >> delete_table_task2
    wait_for_table_replication_task2 >> delete_table_task2
    create_instance_task >> create_table_task >> cluster_update_task
    cluster_update_task >> update_instance_task >> delete_table_task