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