package accesscontrol import ( "context" "errors" "fmt" "testing" "time" "github.com/Tencent/WeKnora/internal/knowledgeversion" "github.com/Tencent/WeKnora/internal/types" "github.com/google/uuid" "github.com/stretchr/testify/require" "gorm.io/gorm" ) type versionCleanupFixture struct { *versionCopyFixture targets []string fail bool after func(string) } func (f *versionCleanupFixture) CleanupKnowledgeVersion(ctx context.Context, target string) error { if !knowledgeversion.IsWorker(ctx) || types.MustTenantIDFromContext(ctx) != 1 { return fmt.Errorf("wrong cleanup workspace") } f.targets = append(f.targets, target) if f.fail { return fmt.Errorf("object store unavailable") } if f.after != nil { f.after(target) } return nil } func newVersionCleanupFixture(t *testing.T) (*Service, *versionCleanupFixture) { t.Helper() s, f, _ := knowledgeVersionFixture(t) require.NoError(t, s.db.AutoMigrate(&types.Tenant{})) require.NoError(t, s.db.Create(&types.Tenant{ID: 1, Name: "cleanup", StorageUsed: 100}).Error) c := &versionCleanupFixture{versionCopyFixture: f} s.SetKnowledgeVersionAdapter(c) return s, c } func addCleanupStage(t *testing.T, s *Service, status string, age time.Duration, temporary bool, state string) KnowledgeVersionJob { t.Helper() id := uuid.NewString() require.NoError(t, s.db.Create(&types.KnowledgeBase{ID: id, TenantID: 1, CreatorID: "alice", Type: "document", IsTemporary: temporary, VersionState: state}).Error) require.NoError(t, s.db.Create(&types.Knowledge{ID: uuid.NewString(), TenantID: 1, KnowledgeBaseID: id, StorageSize: 25}).Error) job := KnowledgeVersionJob{ID: uuid.NewString(), TenantID: 1, OwnerID: "alice", TargetID: id, Status: status, Attempts: 1, AbandonedIDs: types.StringArray{}, UpdatedAt: time.Now().Add(-age)} require.NoError(t, s.db.Create(&job).Error) return job } func TestKnowledgeVersionCleanupDelayAndAccounting(t *testing.T) { s, f := newVersionCleanupFixture(t) recent := addCleanupStage(t, s, "failed", time.Hour, true, "copying") running := addCleanupStage(t, s, "running", 4*time.Hour, true, "copying") snapshot := addCleanupStage(t, s, "failed", 4*time.Hour, true, "snapshot") published := addCleanupStage(t, s, "failed", 4*time.Hour, false, "") old := addCleanupStage(t, s, "failed", 4*time.Hour, true, "copying") require.NoError(t, s.CleanupAbandonedKnowledgeVersions(context.Background())) require.Equal(t, []string{old.TargetID}, f.targets) var count int64 require.NoError(t, s.db.Unscoped().Model(&types.Knowledge{}).Where("knowledge_base_id=?", old.TargetID).Count(&count).Error) require.Zero(t, count) var tenant types.Tenant require.NoError(t, s.db.First(&tenant, 1).Error) require.EqualValues(t, 75, tenant.StorageUsed) require.NoError(t, s.CleanupAbandonedKnowledgeVersions(context.Background())) require.NoError(t, s.db.First(&tenant, 1).Error) require.EqualValues(t, 75, tenant.StorageUsed) for _, job := range []KnowledgeVersionJob{recent, running, snapshot, published} { require.NoError(t, s.db.Model(&types.KnowledgeBase{}).Where("id=?", job.TargetID).Count(&count).Error) require.EqualValues(t, 1, count) } } func TestKnowledgeVersionCleanupRetriesAndRetainsAccountingUntilSuccess(t *testing.T) { s, f := newVersionCleanupFixture(t) job := addCleanupStage(t, s, "failed", 4*time.Hour, true, "copying") f.fail = true require.ErrorContains(t, s.CleanupAbandonedKnowledgeVersions(context.Background()), "object store unavailable") var tenant types.Tenant require.NoError(t, s.db.First(&tenant, 1).Error) require.EqualValues(t, 100, tenant.StorageUsed) f.fail = false require.NoError(t, s.db.Model(&job).UpdateColumn("updated_at", time.Now().Add(-4*time.Hour)).Error) require.NoError(t, s.CleanupAbandonedKnowledgeVersions(context.Background())) require.NoError(t, s.db.First(&tenant, 1).Error) require.EqualValues(t, 75, tenant.StorageUsed) } func TestKnowledgeVersionCleanupPreservesRetryTarget(t *testing.T) { s, f := newVersionCleanupFixture(t) job := addCleanupStage(t, s, "failed", 4*time.Hour, true, "copying") newTarget := uuid.NewString() f.after = func(target string) { require.NoError(t, s.db.Create(&types.KnowledgeBase{ID: newTarget, TenantID: 1, IsTemporary: true, VersionState: "copying"}).Error) require.NoError(t, s.db.Model(&job).Updates(map[string]any{"status": "running", "target_id": newTarget, "abandoned_ids": types.StringArray{target}}).Error) } require.NoError(t, s.CleanupAbandonedKnowledgeVersions(context.Background())) require.NoError(t, s.db.First(&job, "id=?", job.ID).Error) require.Equal(t, newTarget, job.TargetID) require.Empty(t, job.AbandonedIDs) var count int64 require.NoError(t, s.db.Model(&types.KnowledgeBase{}).Where("id=?", newTarget).Count(&count).Error) require.EqualValues(t, 1, count) } func TestKnowledgeVersionCleanupRejectsStorageMismatchAtomically(t *testing.T) { s, _ := newVersionCleanupFixture(t) job := addCleanupStage(t, s, "failed", 4*time.Hour, true, "copying") require.NoError(t, s.db.Model(&types.Tenant{}).Where("id=1").UpdateColumn("storage_used", 0).Error) require.ErrorContains(t, s.CleanupAbandonedKnowledgeVersions(context.Background()), "accounting mismatch") var count int64 require.NoError(t, s.db.Model(&types.Knowledge{}).Where("knowledge_base_id=?", job.TargetID).Count(&count).Error) require.EqualValues(t, 1, count) } func TestKnowledgeVersionIdleQueueIsExpected(t *testing.T) { s, _, _ := knowledgeVersionFixture(t) require.True(t, errors.Is(s.RunKnowledgeVersionJob(context.Background()), gorm.ErrRecordNotFound)) } func TestKnowledgeVersionGraphFingerprintAndScopedCleanup(t *testing.T) { s, _ := newVersionCleanupFixture(t) require.NoError(t, s.db.AutoMigrate(&knowledgeversion.GraphRecord{}, &knowledgeversion.GraphSource{})) job := addCleanupStage(t, s, "failed", 4*time.Hour, true, "copying") for _, kb := range []string{job.TargetID, "published-keep"} { require.NoError(t, s.db.Create(&knowledgeversion.GraphRecord{TenantID: 1, KnowledgeBaseID: kb, SourceTable: "knowledgeentity", LegacyID: 7, Payload: types.JSON(`{"id":7}`)}).Error) require.NoError(t, s.db.Create(&knowledgeversion.GraphSource{TenantID: 1, KnowledgeBaseID: kb, SourceChunkID: "old", Excerpt: "preserve", MappingState: "unmapped"}).Error) } before, err := knowledgeversion.GraphFingerprint(s.db, 1, "published-keep") require.NoError(t, err) require.NotEmpty(t, before) require.NoError(t, s.db.Model(&knowledgeversion.GraphSource{}).Where("knowledge_base_id=?", "published-keep").Update("excerpt", "changed provenance").Error) after, err := knowledgeversion.GraphFingerprint(s.db, 1, "published-keep") require.NoError(t, err) require.NotEqual(t, before, after) require.NoError(t, s.CleanupAbandonedKnowledgeVersions(context.Background())) for _, model := range []any{&knowledgeversion.GraphRecord{}, &knowledgeversion.GraphSource{}} { var count int64 require.NoError(t, s.db.Model(model).Where("knowledge_base_id=?", job.TargetID).Count(&count).Error) require.Zero(t, count) require.NoError(t, s.db.Model(model).Where("knowledge_base_id=?", "published-keep").Count(&count).Error) require.EqualValues(t, 1, count) } }