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)
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)
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)
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,
impersonation_chain=IMPERSONATION_CHAIN,
)
assert not op.poke(None)
mock_hook.assert_called_once_with(
gcp_conn_id=GCP_CONN_ID,
impersonation_chain=IMPERSONATION_CHAIN,
)
def test_wait_not_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(0)
}
op = BigtableTableReplicationCompletedSensor(
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,
)
assert not op.poke(None)
mock_hook.assert_called_once_with(
gcp_conn_id=GCP_CONN_ID,
impersonation_chain=IMPERSONATION_CHAIN,
)
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()
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,
impersonation_chain=IMPERSONATION_CHAIN,
)
err = e.exception
self.assertEqual(str(err), f'Empty parameter: {missing_attribute}')
mock_hook.assert_not_called()
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,