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_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_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)