def test_delete_execute(self, mock_hook): op = BigtableTableDeleteOperator(project_id=PROJECT_ID, instance_id=INSTANCE_ID, table_id=TABLE_ID, task_id="id") op.execute(None) mock_hook.assert_called_once_with() mock_hook.return_value.delete_table.assert_called_once_with( PROJECT_ID, INSTANCE_ID, TABLE_ID)
def test_delete_execute(self, mock_hook): op = BigtableTableDeleteOperator( project_id=PROJECT_ID, instance_id=INSTANCE_ID, table_id=TABLE_ID, task_id="id" ) op.execute(None) mock_hook.assert_called_once_with() mock_hook.return_value.delete_table.assert_called_once_with(PROJECT_ID, INSTANCE_ID, TABLE_ID)
def test_deleting_table_that_doesnt_exists(self, mock_hook): op = BigtableTableDeleteOperator(project_id=PROJECT_ID, instance_id=INSTANCE_ID, table_id=TABLE_ID, task_id="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() mock_hook.return_value.delete_table.assert_called_once_with( PROJECT_ID, INSTANCE_ID, TABLE_ID)
def test_deleting_table_that_doesnt_exists(self, mock_hook): op = BigtableTableDeleteOperator( project_id=PROJECT_ID, instance_id=INSTANCE_ID, table_id=TABLE_ID, task_id="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() mock_hook.return_value.delete_table.assert_called_once_with(PROJECT_ID, INSTANCE_ID, TABLE_ID)
def test_different_error_reraised(self, mock_hook): op = BigtableTableDeleteOperator(project_id=PROJECT_ID, instance_id=INSTANCE_ID, table_id=TABLE_ID, task_id="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() mock_hook.return_value.delete_table.assert_called_once_with( PROJECT_ID, INSTANCE_ID, TABLE_ID)
def test_deleting_table_that_doesnt_exists_empty_project_id( self, mock_hook): op = BigtableTableDeleteOperator(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)
def test_deleting_table_when_instance_doesnt_exists(self, mock_hook): op = BigtableTableDeleteOperator(project_id=PROJECT_ID, instance_id=INSTANCE_ID, table_id=TABLE_ID, task_id="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() mock_hook.return_value.delete_table.assert_not_called()
def test_different_error_reraised(self, mock_hook): op = BigtableTableDeleteOperator( project_id=PROJECT_ID, instance_id=INSTANCE_ID, table_id=TABLE_ID, task_id="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() mock_hook.return_value.delete_table.assert_called_once_with(PROJECT_ID, INSTANCE_ID, TABLE_ID)
def test_deleting_table_when_instance_doesnt_exists(self, mock_hook): op = BigtableTableDeleteOperator( project_id=PROJECT_ID, instance_id=INSTANCE_ID, table_id=TABLE_ID, task_id="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() mock_hook.return_value.delete_table.assert_not_called()
def test_empty_attribute(self, missing_attribute, project_id, instance_id, table_id, mock_hook): with self.assertRaises(AirflowException) as e: BigtableTableDeleteOperator(project_id=project_id, instance_id=instance_id, table_id=table_id, task_id="id") err = e.exception self.assertEqual(str(err), 'Empty parameter: {}'.format(missing_attribute)) mock_hook.assert_not_called()
timeout=180, task_id='wait_for_table_replication_task', ) wait_for_table_replication_task2 = BigtableTableWaitForReplicationSensor( 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 = BigtableTableDeleteOperator( project_id=GCP_PROJECT_ID, instance_id=CBT_INSTANCE_ID, table_id=CBT_TABLE_ID, task_id='delete_table_task', ) delete_table_task2 = BigtableTableDeleteOperator( instance_id=CBT_INSTANCE_ID, table_id=CBT_TABLE_ID, task_id='delete_table_task2', ) # [END howto_operator_gcp_bigtable_table_delete] # pylint: disable=pointless-statement 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 \