WeKnora deployment source with the reviewed wiki identity, summary and recovery fixes

internal/application/service/knowledge_version_cleanup_test.go

@@ -0,0 +1,230 @@
+package service
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"testing"
+
+	"github.com/Tencent/WeKnora/internal/application/repository"
+	filesvc "github.com/Tencent/WeKnora/internal/application/service/file"
+	"github.com/Tencent/WeKnora/internal/knowledgeversion"
+	"github.com/Tencent/WeKnora/internal/types"
+	"github.com/Tencent/WeKnora/internal/types/interfaces"
+	"github.com/stretchr/testify/require"
+)
+
+type versionCleanupEngine struct {
+	parentChildRetrieveEngine
+	deleted []string
+	failure error
+}
+
+func (e *versionCleanupEngine) DeleteByKnowledgeIDList(_ context.Context, ids []string, _ int, _ string) error {
+	if e.failure != nil {
+		return e.failure
+	}
+	e.deleted = append(e.deleted, ids...)
+	return nil
+}
+
+type versionCleanupFiles struct {
+	interfaces.FileService
+	failPath string
+}
+
+func (f *versionCleanupFiles) DeleteFile(ctx context.Context, path string) error {
+	if path == f.failPath {
+		return errors.New("storage unavailable")
+	}
+	return f.FileService.DeleteFile(ctx, path)
+}
+
+func prepareVersionCleanup(t *testing.T, catalog bool) (*versionCopyFixture, *types.Knowledge, []string, *versionCleanupEngine) {
+	t.Helper()
+	f := newVersionCopyFixture(t)
+	f.service.kbService.(versionCopyKBService).bases["target-kb"].IsTemporary = true
+	if catalog {
+		require.NoError(t, f.db.AutoMigrate(&types.StoredResource{}, &types.ResourceBinding{}))
+		f.service.resourceCatalog = &resourceCatalog{repo: repository.NewResourceRepository(f.db)}
+		f.service.fileSvc = filesvc.NewResourceCatalogFileService(f.files, f.service.resourceCatalog)
+	}
+	doc := *f.source
+	doc.ID, doc.KnowledgeBaseID = "target-doc", "target-kb"
+	doc.Metadata, doc.CustomMetadata = nil, nil
+	var err error
+	doc.FilePath, err = f.service.fileSvc.CopyFile(f.ctx, f.raw, 1, doc.ID)
+	require.NoError(t, err)
+	require.NoError(t, f.service.repo.CreateKnowledge(f.ctx, &doc))
+	tracker := &versionTrackingFileService{FileService: f.service.fileSvc, service: f.service, knowledge: &doc}
+	var refs []string
+	for _, name := range []string{"one.png", "two.png"} {
+		ref, err := tracker.SaveBytes(f.ctx, []byte(name), 1, name, false)
+		require.NoError(t, err)
+		refs = append(refs, ref)
+		if catalog {
+			require.NoError(t, f.service.resourceCatalog.Bind(f.ctx, ref, types.ResourceOwnerKnowledge, doc.ID, types.ResourceRelationExtractedImage))
+		}
+	}
+	doc.Metadata = types.JSON(`{"content":"source metadata"}`)
+	require.NoError(t, f.service.preserveKnowledgeVersionInventory(f.ctx, &doc))
+	require.NoError(t, f.service.repo.UpdateKnowledge(f.ctx, &doc))
+	imageInfo, err := json.Marshal([]*types.ImageInfo{{URL: refs[0]}, {URL: refs[1]}})
+	require.NoError(t, err)
+	require.NoError(t, f.service.chunkRepo.CreateChunks(f.ctx, []*types.Chunk{{ID: "target-chunk", TenantID: 1, KnowledgeID: doc.ID, KnowledgeBaseID: doc.KnowledgeBaseID, ChunkType: types.ChunkTypeText, ImageInfo: string(imageInfo)}}))
+	require.NoError(t, f.service.tagRepo.Create(f.ctx, &types.KnowledgeTag{ID: "target-tag", SeqID: 2, TenantID: 1, KnowledgeBaseID: "target-kb", Name: "staged"}))
+	engine := &versionCleanupEngine{}
+	f.service.retrieveEngine = parentChildRetrieveRegistry{engine: engine}
+	return f, &doc, refs, engine
+}
+
+func TestKnowledgeVersionCleanupRejectsUnsafeScope(t *testing.T) {
+	for _, scenario := range []string{"non-worker", "missing-tenant", "cross-tenant", "normal", "snapshot", "not-temporary"} {
+		t.Run(scenario, func(t *testing.T) {
+			f, doc, _, engine := prepareVersionCleanup(t, false)
+			ctx := f.ctx
+			kb := f.service.kbService.(versionCopyKBService).bases["target-kb"]
+			switch scenario {
+			case "non-worker":
+				ctx = context.WithValue(context.Background(), types.TenantIDContextKey, uint64(1))
+			case "missing-tenant":
+				ctx = knowledgeversion.WithWorker(context.Background())
+			case "cross-tenant":
+				kb.TenantID = 2
+			case "normal":
+				kb.VersionState = ""
+			case "snapshot":
+				kb.VersionState = "snapshot"
+			case "not-temporary":
+				kb.IsTemporary = false
+			}
+			require.Error(t, f.service.CleanupKnowledgeVersion(ctx, kb.ID))
+			require.Empty(t, engine.deleted)
+			require.Equal(t, "source raw bytes", readVersionCopyFile(t, f.files, doc.FilePath))
+		})
+	}
+}
+
+func TestKnowledgeVersionCleanupDeletesIndependentObjectsAndIsIdempotent(t *testing.T) {
+	for _, catalog := range []bool{false, true} {
+		t.Run(map[bool]string{false: "legacy", true: "catalog"}[catalog], func(t *testing.T) {
+			f, doc, refs, engine := prepareVersionCleanup(t, catalog)
+			require.NoError(t, f.service.CleanupKnowledgeVersion(f.ctx, "target-kb"))
+			require.NoError(t, f.service.CleanupKnowledgeVersion(f.ctx, "target-kb"))
+			require.Equal(t, []string{doc.ID, doc.ID}, engine.deleted)
+			for _, ref := range append(refs, doc.FilePath) {
+				_, err := f.service.fileSvc.GetFile(f.ctx, ref)
+				require.Error(t, err)
+			}
+			require.Equal(t, "source raw bytes", readVersionCopyFile(t, f.files, f.raw))
+			require.Equal(t, "chunk image bytes", readVersionCopyFile(t, f.files, f.image))
+			var count int64
+			require.NoError(t, f.db.Unscoped().Model(&types.Knowledge{}).Where("id = ?", doc.ID).Count(&count).Error)
+			require.EqualValues(t, 1, count)
+			require.NoError(t, f.db.Model(&types.Chunk{}).Where("knowledge_base_id = ?", "target-kb").Count(&count).Error)
+			require.Zero(t, count)
+			tags, _, err := f.service.tagRepo.ListByKB(f.ctx, 1, "target-kb", &types.Pagination{Page: 1, PageSize: 100}, "")
+			require.NoError(t, err)
+			require.Empty(t, tags)
+		})
+	}
+}
+
+func TestKnowledgeVersionCleanupRetriesStorageAndIndexFailures(t *testing.T) {
+	f, doc, refs, engine := prepareVersionCleanup(t, false)
+	engine.failure = errors.New("vector unavailable")
+	require.ErrorContains(t, f.service.CleanupKnowledgeVersion(f.ctx, "target-kb"), "vector unavailable")
+	require.Equal(t, "one.png", readVersionCopyFile(t, f.files, refs[0]))
+	engine.failure = nil
+	failing := &versionCleanupFiles{FileService: f.service.fileSvc, failPath: refs[1]}
+	f.service.fileSvc = failing
+	require.ErrorContains(t, f.service.CleanupKnowledgeVersion(f.ctx, "target-kb"), "storage unavailable")
+	retained, err := f.service.repo.GetKnowledgeByID(f.ctx, 1, doc.ID)
+	require.NoError(t, err)
+	_, inventory, err := readKnowledgeVersionInventory(retained.Metadata, doc.ID)
+	require.NoError(t, err)
+	require.True(t, inventory.Objects[0].Done)
+	require.False(t, inventory.Objects[1].Done)
+	failing.failPath = ""
+	require.NoError(t, f.service.CleanupKnowledgeVersion(f.ctx, "target-kb"))
+}
+
+func TestKnowledgeVersionCleanupPartialDocumentRetainsImageJournal(t *testing.T) {
+	f, doc, refs, _ := prepareVersionCleanup(t, false)
+	require.NoError(t, f.service.chunkRepo.DeleteChunksByKnowledgeID(f.ctx, 1, doc.ID))
+	require.NoError(t, f.service.CleanupKnowledgeVersion(f.ctx, "target-kb"))
+	for _, ref := range refs {
+		_, err := f.files.GetFile(f.ctx, ref)
+		require.Error(t, err)
+	}
+}
+
+func TestKnowledgeVersionCleanupRetainsSharedResource(t *testing.T) {
+	f, _, refs, _ := prepareVersionCleanup(t, true)
+	require.NoError(t, f.service.resourceCatalog.Bind(f.ctx, refs[0], types.ResourceOwnerMessage, "message-owner", types.ResourceRelationAttachment))
+	require.NoError(t, f.service.CleanupKnowledgeVersion(f.ctx, "target-kb"))
+	require.Equal(t, "one.png", readVersionCopyFile(t, f.service.fileSvc, refs[0]))
+}
+
+func TestKnowledgeVersionCleanupRefusesSourceFileAndUntrackedImage(t *testing.T) {
+	for _, scenario := range []string{"source-file", "untracked-image"} {
+		t.Run(scenario, func(t *testing.T) {
+			f, doc, _, engine := prepareVersionCleanup(t, false)
+			if scenario == "source-file" {
+				require.NoError(t, f.service.repo.UpdateKnowledgeColumns(f.ctx, doc.ID, map[string]interface{}{"file_path": f.raw}))
+			} else {
+				info, err := json.Marshal([]*types.ImageInfo{{URL: f.image}})
+				require.NoError(t, err)
+				require.NoError(t, f.db.Model(&types.Chunk{}).Where("id = ?", "target-chunk").Update("image_info", string(info)).Error)
+			}
+			require.Error(t, f.service.CleanupKnowledgeVersion(f.ctx, "target-kb"))
+			require.Empty(t, engine.deleted)
+			require.Equal(t, "source raw bytes", readVersionCopyFile(t, f.files, f.raw))
+			require.Equal(t, "chunk image bytes", readVersionCopyFile(t, f.files, f.image))
+		})
+	}
+}
+
+type versionCleanupReleaseFailure struct{ interfaces.ResourceCatalog }
+
+func (versionCleanupReleaseFailure) Release(context.Context, string, string, string) (int64, error) {
+	return -1, errors.New("binding lookup failed")
+}
+
+func TestKnowledgeVersionCleanupDoesNotDeleteOnReleaseFailure(t *testing.T) {
+	f, doc, refs, _ := prepareVersionCleanup(t, true)
+	nativeCatalog := f.service.resourceCatalog
+	f.service.resourceCatalog = versionCleanupReleaseFailure{nativeCatalog}
+	require.ErrorContains(t, f.service.CleanupKnowledgeVersion(f.ctx, "target-kb"), "binding lookup failed")
+	require.Equal(t, "one.png", readVersionCopyFile(t, f.service.fileSvc, refs[0]))
+	_, err := f.service.repo.GetKnowledgeByID(f.ctx, 1, doc.ID)
+	require.NoError(t, err)
+	f.service.resourceCatalog = nativeCatalog
+	require.NoError(t, f.service.CleanupKnowledgeVersion(f.ctx, "target-kb"))
+}
+
+type versionCleanupJournalFailure struct {
+	interfaces.KnowledgeRepository
+	writes int
+}
+
+func (r *versionCleanupJournalFailure) UpdateKnowledgeColumns(ctx context.Context, id string, values map[string]interface{}) error {
+	r.writes++
+	if r.writes == 2 {
+		return errors.New("journal update failed after file removal")
+	}
+	return r.KnowledgeRepository.UpdateKnowledgeColumns(ctx, id, values)
+}
+
+func TestKnowledgeVersionCleanupRetriesAfterFileDeletedBeforeJournalUpdate(t *testing.T) {
+	for _, catalog := range []bool{false, true} {
+		t.Run(map[bool]string{false: "legacy", true: "catalog"}[catalog], func(t *testing.T) {
+			f, doc, _, _ := prepareVersionCleanup(t, catalog)
+			f.service.repo = &versionCleanupJournalFailure{KnowledgeRepository: f.service.repo}
+			require.ErrorContains(t, f.service.CleanupKnowledgeVersion(f.ctx, "target-kb"), "journal update failed")
+			_, err := f.service.repo.GetKnowledgeByID(f.ctx, 1, doc.ID)
+			require.NoError(t, err)
+			require.NoError(t, f.service.CleanupKnowledgeVersion(f.ctx, "target-kb"))
+		})
+	}
+}