def test_creating_table_that_exists_with_different_column_families_gc_rule_in_the_table( self, mock_hook): op = BigtableTableCreateOperator( project_id=PROJECT_ID, instance_id=INSTANCE_ID, table_id=TABLE_ID, initial_split_keys=INITIAL_SPLIT_KEYS, column_families={"cf-id": MaxVersionsGCRule(1)}, task_id="id") cf_mock = mock.Mock() cf_mock.gc_rule = mock.Mock(return_value=MaxVersionsGCRule(2)) mock_hook.return_value.get_column_families_for_table.return_value = { "cf-id": cf_mock } mock_hook.return_value.create_table.side_effect = mock.Mock( side_effect=google.api_core.exceptions.AlreadyExists( "Table already exists.")) with self.assertRaises(AirflowException) as e: op.execute(None) err = e.exception self.assertEqual( str(err), "Table '{}' already exists with different Column Families.".format( TABLE_ID)) mock_hook.assert_called_once_with()
def test_creating_table_that_exists_with_different_column_families_ids_in_the_table( self, mock_hook): op = BigtableTableCreateOperator(project_id=PROJECT_ID, instance_id=INSTANCE_ID, table_id=TABLE_ID, initial_split_keys=INITIAL_SPLIT_KEYS, column_families=EMPTY_COLUMN_FAMILIES, task_id="id", gcp_conn_id=GCP_CONN_ID) mock_hook.return_value.get_column_families_for_table.return_value = { "existing_family": None } mock_hook.return_value.create_table.side_effect = mock.Mock( side_effect=google.api_core.exceptions.AlreadyExists( "Table already exists.")) with self.assertRaises(AirflowException) as e: op.execute(None) err = e.exception self.assertEqual( str(err), "Table '{}' already exists with different Column Families.".format( TABLE_ID)) mock_hook.assert_called_once_with(gcp_conn_id=GCP_CONN_ID)
def test_creating_table_that_exists_with_different_column_families_gc_rule_in__table( self, mock_hook): op = BigtableTableCreateOperator( project_id=PROJECT_ID, instance_id=INSTANCE_ID, table_id=TABLE_ID, initial_split_keys=INITIAL_SPLIT_KEYS, column_families={"cf-id": MaxVersionsGCRule(1)}, task_id="id" ) cf_mock = mock.Mock() cf_mock.gc_rule = mock.Mock(return_value=MaxVersionsGCRule(2)) mock_hook.return_value.get_column_families_for_table.return_value = { "cf-id": cf_mock } mock_hook.return_value.create_table.side_effect = mock.Mock( side_effect=google.api_core.exceptions.AlreadyExists("Table already exists.")) with self.assertRaises(AirflowException) as e: op.execute(None) err = e.exception self.assertEqual( str(err), "Table '{}' already exists with different Column Families.".format(TABLE_ID) ) mock_hook.assert_called_once_with()
def test_create_execute(self, mock_hook): op = BigtableTableCreateOperator(project_id=PROJECT_ID, instance_id=INSTANCE_ID, table_id=TABLE_ID, initial_split_keys=INITIAL_SPLIT_KEYS, column_families=EMPTY_COLUMN_FAMILIES, task_id="id") instance = mock_hook.return_value.get_instance.return_value = mock.Mock( Instance) op.execute(None) mock_hook.assert_called_once_with() mock_hook.return_value.create_table.assert_called_once_with( instance, TABLE_ID, INITIAL_SPLIT_KEYS, EMPTY_COLUMN_FAMILIES)
def test_create_execute(self, mock_hook): op = BigtableTableCreateOperator( project_id=PROJECT_ID, instance_id=INSTANCE_ID, table_id=TABLE_ID, initial_split_keys=INITIAL_SPLIT_KEYS, column_families=EMPTY_COLUMN_FAMILIES, task_id="id" ) instance = mock_hook.return_value.get_instance.return_value = mock.Mock(Instance) op.execute(None) mock_hook.assert_called_once_with() mock_hook.return_value.create_table.assert_called_once_with( instance, TABLE_ID, INITIAL_SPLIT_KEYS, EMPTY_COLUMN_FAMILIES)
def test_instance_not_exists(self, mock_hook): op = BigtableTableCreateOperator(project_id=PROJECT_ID, instance_id=INSTANCE_ID, table_id=TABLE_ID, initial_split_keys=INITIAL_SPLIT_KEYS, column_families=EMPTY_COLUMN_FAMILIES, 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 in project '{}'.".format( INSTANCE_ID, PROJECT_ID)) mock_hook.assert_called_once_with()
def test_instance_not_exists(self, mock_hook): op = BigtableTableCreateOperator( project_id=PROJECT_ID, instance_id=INSTANCE_ID, table_id=TABLE_ID, initial_split_keys=INITIAL_SPLIT_KEYS, column_families=EMPTY_COLUMN_FAMILIES, 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 in project '{}'.".format(INSTANCE_ID, PROJECT_ID) ) mock_hook.assert_called_once_with()
def test_creating_table_that_exists(self, mock_hook): op = BigtableTableCreateOperator(project_id=PROJECT_ID, instance_id=INSTANCE_ID, table_id=TABLE_ID, initial_split_keys=INITIAL_SPLIT_KEYS, column_families=EMPTY_COLUMN_FAMILIES, task_id="id") mock_hook.return_value.get_column_families_for_table.return_value = EMPTY_COLUMN_FAMILIES instance = mock_hook.return_value.get_instance.return_value = mock.Mock( Instance) mock_hook.return_value.create_table.side_effect = mock.Mock( side_effect=google.api_core.exceptions.AlreadyExists( "Table already exists.")) op.execute(None) mock_hook.assert_called_once_with() mock_hook.return_value.create_table.assert_called_once_with( instance, TABLE_ID, INITIAL_SPLIT_KEYS, EMPTY_COLUMN_FAMILIES)
def test_creating_table_that_exists(self, mock_hook): op = BigtableTableCreateOperator( project_id=PROJECT_ID, instance_id=INSTANCE_ID, table_id=TABLE_ID, initial_split_keys=INITIAL_SPLIT_KEYS, column_families=EMPTY_COLUMN_FAMILIES, task_id="id" ) mock_hook.return_value.get_column_families_for_table.return_value = EMPTY_COLUMN_FAMILIES instance = mock_hook.return_value.get_instance.return_value = mock.Mock(Instance) mock_hook.return_value.create_table.side_effect = mock.Mock( side_effect=google.api_core.exceptions.AlreadyExists("Table already exists.")) op.execute(None) mock_hook.assert_called_once_with() mock_hook.return_value.create_table.assert_called_once_with( instance, TABLE_ID, INITIAL_SPLIT_KEYS, EMPTY_COLUMN_FAMILIES)
def test_empty_attribute(self, missing_attribute, project_id, instance_id, table_id, mock_hook): with self.assertRaises(AirflowException) as e: BigtableTableCreateOperator(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()
# [START howto_operator_gcp_bigtable_instance_delete] delete_instance_task = BigtableInstanceDeleteOperator( project_id=GCP_PROJECT_ID, instance_id=CBT_INSTANCE_ID, task_id='delete_instance_task', ) delete_instance_task2 = BigtableInstanceDeleteOperator( instance_id=CBT_INSTANCE_ID, task_id='delete_instance_task2', ) # [END howto_operator_gcp_bigtable_instance_delete] # [START howto_operator_gcp_bigtable_table_create] create_table_task = BigtableTableCreateOperator( project_id=GCP_PROJECT_ID, instance_id=CBT_INSTANCE_ID, table_id=CBT_TABLE_ID, task_id='create_table', ) create_table_task2 = BigtableTableCreateOperator( instance_id=CBT_INSTANCE_ID, table_id=CBT_TABLE_ID, task_id='create_table_task2', ) create_table_task >> create_table_task2 # pylint: disable=pointless-statement # [END howto_operator_gcp_bigtable_table_create] # [START howto_operator_gcp_bigtable_table_wait_for_replication] wait_for_table_replication_task = BigtableTableWaitForReplicationSensor( project_id=GCP_PROJECT_ID, instance_id=CBT_INSTANCE_ID, table_id=CBT_TABLE_ID,