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

internal/accesscontrol/resource_version_knowledge_cleanup_test.go

@@ -0,0 +1,152 @@
+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)
+	}
+}