xinmi/internal/accesscontrol/resource_version_knowledge_cleanup_test.go

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