package container import ( "context" "database/sql" "testing" "time" "github.com/Tencent/WeKnora/internal/types" "github.com/stretchr/testify/require" "gorm.io/driver/sqlite" "gorm.io/gorm" ) const pendingOpsDDL = ` CREATE TABLE IF NOT EXISTS task_pending_ops ( id INTEGER PRIMARY KEY AUTOINCREMENT, tenant_id INTEGER NOT NULL DEFAULT 0, task_type VARCHAR(64) NOT NULL DEFAULT '', scope VARCHAR(32) NOT NULL DEFAULT '', scope_id VARCHAR(64) NOT NULL DEFAULT '', op VARCHAR(32) NOT NULL DEFAULT '', dedup_key VARCHAR(128) NOT NULL DEFAULT '', payload TEXT, fail_count INTEGER NOT NULL DEFAULT 0, enqueued_at DATETIME, claimed_at DATETIME ); ` func openPendingOpsDB(t *testing.T) *gorm.DB { t.Helper() db, err := gorm.Open(sqlite.Open("file::memory:?cache=shared&_uniq="+t.Name()), &gorm.Config{}) require.NoError(t, err) require.NoError(t, db.Exec(pendingOpsDDL).Error) t.Cleanup(func() { if sqlDB, err := db.DB(); err == nil { _ = sqlDB.Close() } }) return db } // claimOf reads the claim column as nullable, so a released claim is distinguishable // from a timestamp. func claimOf(t *testing.T, db *gorm.DB, id int64) (time.Time, bool) { t.Helper() var got sql.NullTime require.NoError(t, db.Table("task_pending_ops").Select("claimed_at").Where("id = ?", id).Scan(&got).Error) return got.Time, got.Valid } func seedPendingOp(t *testing.T, db *gorm.DB, taskType, scopeID string, claimedAt *time.Time) int64 { t.Helper() row := types.TaskPendingOp{ TenantID: 10059, TaskType: taskType, Scope: types.TaskScopeKnowledgeBase, ScopeID: scopeID, Op: "ingest", EnqueuedAt: time.Now().UTC(), ClaimedAt: claimedAt, } require.NoError(t, db.Create(&row).Error) return row.ID } // A claim stamped by a worker that has since died must be released at startup, so a // restarted deployment resumes instead of waiting out the 90 minute stale threshold. func TestSweepOrphanedWikiClaimsReleasesStaleClaimsOnly(t *testing.T) { db := openPendingOpsDB(t) now := time.Now().UTC() stale := now.Add(-20 * time.Minute) fresh := now.Add(-1 * time.Minute) nonWikiStale := now.Add(-45 * time.Minute) staleID := seedPendingOp(t, db, types.TypeWikiIngest, "kb-stale", &stale) freshID := seedPendingOp(t, db, types.TypeWikiIngest, "kb-fresh", &fresh) unclaimedID := seedPendingOp(t, db, types.TypeWikiIngest, "kb-unclaimed", nil) nonWikiID := seedPendingOp(t, db, "some:other:task", "kb-other", &nonWikiStale) released, err := sweepOrphanedWikiClaims(context.Background(), db, now) require.NoError(t, err) require.EqualValues(t, 1, released) _, staleStillClaimed := claimOf(t, db, staleID) require.False(t, staleStillClaimed, "the stale wiki claim is released") _, freshStillClaimed := claimOf(t, db, freshID) require.True(t, freshStillClaimed, "a live worker's fresh claim is untouched") _, unclaimedStillClaimed := claimOf(t, db, unclaimedID) require.False(t, unclaimedStillClaimed) _, nonWikiStillClaimed := claimOf(t, db, nonWikiID) require.True(t, nonWikiStillClaimed, "non-wiki queues keep their own claim policy") } // Without the durable table the sweep is a no-op rather than an error, because the // Lite deployment may run without the task queue at all. func TestSweepOrphanedWikiClaimsWithoutTableIsNoop(t *testing.T) { db, err := gorm.Open(sqlite.Open("file::memory:?cache=shared&_uniq=notable"), &gorm.Config{}) require.NoError(t, err) released, err := sweepOrphanedWikiClaims(context.Background(), db, time.Now()) require.NoError(t, err) require.EqualValues(t, 0, released) }