xinmi/internal/container/recover_pending_wiki_tasks_test.go

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