231 lines
9.8 KiB
Go
231 lines
9.8 KiB
Go
|
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"))
|
||
|
|
})
|
||
|
|
}
|
||
|
|
}
|