Exemple #1
0
    def test_wait_no_instance(self, mock_hook):
        mock_hook.return_value.get_instance.return_value = None

        op = BigtableTableReplicationCompletedSensor(project_id=PROJECT_ID,
                                                     instance_id=INSTANCE_ID,
                                                     table_id=TABLE_ID,
                                                     task_id="id",
                                                     gcp_conn_id=GCP_CONN_ID)
        self.assertFalse(op.poke(None))
        mock_hook.assert_called_once_with(gcp_conn_id=GCP_CONN_ID)
Exemple #2
0
 def test_wait_ready(self, mock_hook):
     mock_hook.return_value.get_instance.return_value = mock.Mock(Instance)
     mock_hook.return_value.get_cluster_states_for_table.return_value = {
         "cl-id": ClusterState(4)
     }
     op = BigtableTableReplicationCompletedSensor(project_id=PROJECT_ID,
                                                  instance_id=INSTANCE_ID,
                                                  table_id=TABLE_ID,
                                                  task_id="id",
                                                  gcp_conn_id=GCP_CONN_ID)
     self.assertTrue(op.poke(None))
     mock_hook.assert_called_once_with(gcp_conn_id=GCP_CONN_ID)
Exemple #3
0
    def test_wait_no_table(self, mock_hook):
        mock_hook.return_value.get_instance.return_value = mock.Mock(Instance)
        mock_hook.return_value.get_cluster_states_for_table.side_effect = mock.Mock(
            side_effect=google.api_core.exceptions.NotFound(
                "Table not found."))

        op = BigtableTableReplicationCompletedSensor(project_id=PROJECT_ID,
                                                     instance_id=INSTANCE_ID,
                                                     table_id=TABLE_ID,
                                                     task_id="id",
                                                     gcp_conn_id=GCP_CONN_ID)
        self.assertFalse(op.poke(None))
        mock_hook.assert_called_once_with(gcp_conn_id=GCP_CONN_ID)
Exemple #4
0
 def test_empty_attribute(self, missing_attribute, project_id, instance_id,
                          table_id, mock_hook):
     with self.assertRaises(AirflowException) as e:
         BigtableTableReplicationCompletedSensor(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()
Exemple #5
0
        table_id=CBT_TABLE_ID,
        task_id='create_table',
    )
    create_table_task2 = BigtableCreateTableOperator(
        instance_id=CBT_INSTANCE_ID,
        table_id=CBT_TABLE_ID,
        task_id='create_table_task2',
    )
    create_table_task >> create_table_task2
    # [END howto_operator_gcp_bigtable_table_create]

    # [START howto_operator_gcp_bigtable_table_wait_for_replication]
    wait_for_table_replication_task = BigtableTableReplicationCompletedSensor(
        project_id=GCP_PROJECT_ID,
        instance_id=CBT_INSTANCE_ID,
        table_id=CBT_TABLE_ID,
        poke_interval=int(CBT_POKE_INTERVAL),
        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,