xinmi/internal/application/repository/knowledgebase_delete_test.go

package repository

import (
	"context"
	"testing"
	"time"

	"github.com/Tencent/WeKnora/internal/types"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"gorm.io/gorm"
)

func setupKBDeleteTestDB(t *testing.T) *gorm.DB {
	t.Helper()
	db := setupKBTestDB(t)
	require.NoError(t, db.Exec(knowledgeTagTestDDL).Error)
	sqlDB, err := db.DB()
	require.NoError(t, err)
	t.Cleanup(func() { _ = sqlDB.Close() })
	return db
}

func kbDeleteContext(tenantID uint64) context.Context {
	return context.WithValue(context.Background(), types.TenantIDContextKey, tenantID)
}

func seedKBDeleteTag(t *testing.T, db *gorm.DB, kb *types.KnowledgeBase, tagID string) {
	t.Helper()
	require.NoError(t, db.Create(&types.KnowledgeTag{
		ID: tagID, TenantID: kb.TenantID, KnowledgeBaseID: kb.ID, Name: tagID,
	}).Error)
}

func assertKBDeleteTableCount(t *testing.T, db *gorm.DB, table string, want int64) {
	t.Helper()
	var got int64
	require.NoError(t, db.Table(table).Count(&got).Error)
	assert.Equal(t, want, got, table)
}

func TestKnowledgeBaseDeleteCleansEmptyCategoriesAndIsIdempotent(t *testing.T) {
	db := setupKBDeleteTestDB(t)
	kb := makeKB(nil)
	require.NoError(t, db.Create(kb).Error)
	seedKBDeleteTag(t, db, kb, "unused-category")
	repo := &knowledgeBaseRepository{db: db}
	ctx := kbDeleteContext(kb.TenantID)

	require.NoError(t, repo.DeleteKnowledgeBase(ctx, kb.ID))
	var deleted types.KnowledgeBase
	require.NoError(t, db.Unscoped().First(&deleted, "id = ?", kb.ID).Error)
	require.True(t, deleted.DeletedAt.Valid)
	assertKBDeleteTableCount(t, db, "knowledge_tags", 0)
	require.NoError(t, repo.DeleteKnowledgeBase(ctx, kb.ID))
	require.NoError(t, repo.CleanupDeletedKnowledgeBaseCategories(ctx, kb.ID))
	var retried types.KnowledgeBase
	require.NoError(t, db.Unscoped().First(&retried, "id = ?", kb.ID).Error)
	assert.Equal(t, deleted.DeletedAt, retried.DeletedAt)
}

func TestKnowledgeBaseDeleteCleansRelationsAndPreservesOtherResources(t *testing.T) {
	db := setupKBDeleteTestDB(t)
	target, sibling, foreign := makeKB(nil), makeKB(nil), makeKB(nil)
	foreign.TenantID = 2
	for i, kb := range []*types.KnowledgeBase{target, sibling, foreign} {
		require.NoError(t, db.Create(kb).Error)
		seedKBDeleteTag(t, db, kb, []string{"target-tag", "sibling-tag", "foreign-tag"}[i])
	}
	for _, doc := range []struct {
		id string
		kb *types.KnowledgeBase
	}{
		{"target-live", target}, {"target-deleted", target}, {"sibling-doc", sibling}, {"foreign-doc", foreign},
	} {
		require.NoError(t, db.Exec("INSERT INTO knowledges (id, tenant_id, knowledge_base_id) VALUES (?, ?, ?)",
			doc.id, doc.kb.TenantID, doc.kb.ID).Error)
	}
	require.NoError(t, db.Exec("UPDATE knowledges SET deleted_at = ? WHERE id = ?", time.Now(), "target-deleted").Error)
	for _, rel := range []types.KnowledgeTagRelation{
		{KnowledgeID: "target-live", TagID: "target-tag"},
		{KnowledgeID: "target-deleted", TagID: "target-tag"},
		// Repair partial old cleanup even when its category or document is already absent.
		{KnowledgeID: "target-deleted", TagID: "already-removed-tag"},
		{KnowledgeID: "already-removed-doc", TagID: "target-tag"},
		{KnowledgeID: "sibling-doc", TagID: "sibling-tag"},
		{KnowledgeID: "foreign-doc", TagID: "foreign-tag"},
	} {
		require.NoError(t, db.Create(&rel).Error)
	}
	require.NoError(t, db.Exec("INSERT INTO chunks (id, tenant_id, knowledge_base_id, tag_id) VALUES (?, ?, ?, ?)",
		"target-chunk", target.TenantID, target.ID, "target-tag").Error)

	repo := &knowledgeBaseRepository{db: db}
	require.NoError(t, repo.DeleteKnowledgeBase(kbDeleteContext(1), target.ID))
	assertKBDeleteTableCount(t, db, "knowledge_tags", 2)
	assertKBDeleteTableCount(t, db, "knowledges", 4)
	assertKBDeleteTableCount(t, db, "chunks", 1)
	var relations []types.KnowledgeTagRelation
	require.NoError(t, db.Order("knowledge_id").Find(&relations).Error)
	require.Len(t, relations, 2)
	assert.Equal(t, "foreign-doc", relations[0].KnowledgeID)
	assert.Equal(t, "foreign-tag", relations[0].TagID)
	assert.Equal(t, "sibling-doc", relations[1].KnowledgeID)
	assert.Equal(t, "sibling-tag", relations[1].TagID)
	for _, kb := range []*types.KnowledgeBase{sibling, foreign} {
		got := reloadKB(t, db, kb.ID)
		assert.False(t, got.DeletedAt.Valid)
	}
}

func TestKnowledgeBaseDeleteRejectsUnscopedOrWrongTenant(t *testing.T) {
	tests := []struct {
		name    string
		ctx     context.Context
		emptyID bool
	}{
		{"missing tenant", context.Background(), false},
		{"zero tenant", kbDeleteContext(0), false},
		{"wrong tenant", kbDeleteContext(2), false},
		{"empty ID", kbDeleteContext(1), true},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			db := setupKBDeleteTestDB(t)
			kb := makeKB(nil)
			require.NoError(t, db.Create(kb).Error)
			seedKBDeleteTag(t, db, kb, "retained-tag")
			repo := &knowledgeBaseRepository{db: db}
			id := kb.ID
			if tt.emptyID {
				id = ""
			}
			require.Error(t, repo.DeleteKnowledgeBase(tt.ctx, id))
			require.Error(t, repo.CleanupDeletedKnowledgeBaseCategories(tt.ctx, id))
			assert.False(t, reloadKB(t, db, kb.ID).DeletedAt.Valid)
			assertKBDeleteTableCount(t, db, "knowledge_tags", 1)
		})
	}
}

func TestKnowledgeBaseCategoryCleanupRequiresDeletedKB(t *testing.T) {
	db := setupKBDeleteTestDB(t)
	kb := makeKB(nil)
	require.NoError(t, db.Create(kb).Error)
	seedKBDeleteTag(t, db, kb, "retained-tag")
	repo := &knowledgeBaseRepository{db: db}
	ctx := kbDeleteContext(1)
	require.ErrorContains(t, repo.CleanupDeletedKnowledgeBaseCategories(ctx, kb.ID), "requires a deleted")
	assertKBDeleteTableCount(t, db, "knowledge_tags", 1)
	require.NoError(t, db.Delete(kb).Error)
	require.NoError(t, repo.CleanupDeletedKnowledgeBaseCategories(ctx, kb.ID))
	assertKBDeleteTableCount(t, db, "knowledge_tags", 0)
	require.ErrorIs(t, repo.CleanupDeletedKnowledgeBaseCategories(ctx, "missing-kb"), ErrKnowledgeBaseNotFound)
}

func TestKnowledgeBaseDeleteRollsBackCategoryAndKBChanges(t *testing.T) {
	for _, trigger := range []struct{ name, sql string }{
		{"category failure", "CREATE TRIGGER fail_delete BEFORE DELETE ON knowledge_tags BEGIN SELECT RAISE(ABORT, 'category failed'); END"},
		{"KB failure", "CREATE TRIGGER fail_delete BEFORE UPDATE OF deleted_at ON knowledge_bases BEGIN SELECT RAISE(ABORT, 'KB failed'); END"},
	} {
		t.Run(trigger.name, func(t *testing.T) {
			db := setupKBDeleteTestDB(t)
			kb := makeKB(nil)
			require.NoError(t, db.Create(kb).Error)
			seedKBDeleteTag(t, db, kb, "retained-tag")
			require.NoError(t, db.Create(&types.KnowledgeTagRelation{KnowledgeID: "doc", TagID: "retained-tag"}).Error)
			require.NoError(t, db.Exec(trigger.sql).Error)
			repo := &knowledgeBaseRepository{db: db}
			require.Error(t, repo.DeleteKnowledgeBase(kbDeleteContext(1), kb.ID))
			assert.False(t, reloadKB(t, db, kb.ID).DeletedAt.Valid)
			assertKBDeleteTableCount(t, db, "knowledge_tags", 1)
			assertKBDeleteTableCount(t, db, "knowledge_tag_relations", 1)
		})
	}
}

func TestKnowledgeBaseDeleteRejectsLateCategoryWrites(t *testing.T) {
	db := setupKBDeleteTestDB(t)
	kb := makeKB(nil)
	require.NoError(t, db.Create(kb).Error)
	ctx := kbDeleteContext(1)
	tagRepo := &knowledgeTagRepository{db: db}
	knowledgeRepo := &knowledgeRepository{db: db}
	tag := &types.KnowledgeTag{ID: "category", TenantID: 1, KnowledgeBaseID: kb.ID, Name: "original"}
	require.NoError(t, tagRepo.Create(ctx, tag))
	tag.Name = "updated"
	require.NoError(t, tagRepo.Update(ctx, tag))
	require.NoError(t, db.Exec("INSERT INTO knowledges (id, tenant_id, knowledge_base_id) VALUES (?, ?, ?)", "doc", 1, kb.ID).Error)
	require.NoError(t, knowledgeRepo.SetKnowledgeTags(ctx, "doc", []string{tag.ID}))
	kbRepo := &knowledgeBaseRepository{db: db}
	require.NoError(t, kbRepo.DeleteKnowledgeBase(ctx, kb.ID))

	require.ErrorIs(t, tagRepo.Update(ctx, tag), ErrKnowledgeBaseNotFound)
	require.ErrorIs(t, tagRepo.Create(ctx, &types.KnowledgeTag{ID: "late", TenantID: 1, KnowledgeBaseID: kb.ID, Name: "late"}), ErrKnowledgeBaseNotFound)
	require.ErrorIs(t, knowledgeRepo.SetKnowledgeTags(ctx, "doc", []string{tag.ID}), ErrKnowledgeBaseNotFound)
	require.ErrorIs(t, knowledgeRepo.AddKnowledgeTagRelations(ctx, 1, kb.ID, "doc", []string{tag.ID}), ErrKnowledgeBaseNotFound)
	assertKBDeleteTableCount(t, db, "knowledge_tags", 0)
	assertKBDeleteTableCount(t, db, "knowledge_tag_relations", 0)
}