xinmi/internal/application/repository/memory_recovery_test.go

package repository

import (
	"context"
	"fmt"
	"os"
	"sync"
	"testing"
	"time"

	"github.com/Tencent/WeKnora/internal/types"
	"github.com/Tencent/WeKnora/internal/types/interfaces"
	"github.com/stretchr/testify/require"
	"gorm.io/driver/sqlite"
	"gorm.io/gorm"
)

func TestMemoryRecoveryReservationsAndExecutionAreExclusive(t *testing.T) {
	db, err := gorm.Open(sqlite.Open("file:"+t.Name()+"?mode=memory&cache=shared"), &gorm.Config{})
	require.NoError(t, err)
	sqlDB, err := db.DB()
	require.NoError(t, err)
	sqlDB.SetMaxOpenConns(1)
	require.NoError(t, db.AutoMigrate(&types.MemorySubject{}))
	repo := NewMemoryRepository(db)
	ctx := context.Background()
	scope := interfaces.MemoryScope{TenantID: 1, SubjectID: "web_user:alice"}
	_, err = repo.EnsureSubject(ctx, scope)
	require.NoError(t, err)
	_, _, err = repo.EnqueuePendingSession(ctx, scope, "s", time.Minute)
	require.NoError(t, err)
	var wg sync.WaitGroup
	results := make(chan *types.MemoryExtractPayload, 12)
	errors := make(chan error, 12)
	for range 12 {
		wg.Add(1)
		go func() {
			defer wg.Done()
			payload, err := repo.ReserveExtraction(ctx, scope, types.MemoryExtractPayload{}, 0, time.Minute)
			results <- payload
			errors <- err
		}()
	}
	wg.Wait()
	close(results)
	close(errors)
	for err := range errors {
		require.NoError(t, err)
	}
	var winner *types.MemoryExtractPayload
	for payload := range results {
		if payload != nil {
			require.Nil(t, winner)
			winner = payload
		}
	}
	require.NotNil(t, winner)
	for i := range 2 {
		started, err := repo.StartExtraction(ctx, scope, winner.AttemptID, time.Minute)
		require.NoError(t, err)
		require.Equal(t, i == 0, started)
	}
	for failures := 1; failures <= types.MaxMemoryExtractionFailures; failures++ {
		require.NoError(t, repo.FailExtraction(ctx, scope, winner.AttemptID))
		subject, err := repo.GetSubject(ctx, scope)
		require.NoError(t, err)
		require.Equal(t, failures, subject.ExtractFailures)
		if failures == types.MaxMemoryExtractionFailures {
			require.Nil(t, subject.ExtractNextAt)
			winner, err = repo.ReserveExtraction(ctx, scope, types.MemoryExtractPayload{}, 0, time.Minute)
			require.NoError(t, err)
			require.Nil(t, winner)
			break
		}
		require.InDelta(t, extractionRetryDelay(failures).Seconds(), time.Until(*subject.ExtractNextAt).Seconds(), 1)
		require.LessOrEqual(t, time.Until(*subject.ExtractNextAt), 6*time.Hour)
		require.NoError(t, repo.DeferExtraction(ctx, scope, -time.Second))
		winner, err = repo.ReserveExtraction(ctx, scope, types.MemoryExtractPayload{}, 0, time.Minute)
		require.NoError(t, err)
		require.NotNil(t, winner)
	}
}

func TestMemoryRecoveryScanIsBoundedAndExcludesIdleScopes(t *testing.T) {
	db, err := gorm.Open(sqlite.Open("file:"+t.Name()+"?mode=memory&cache=shared"), &gorm.Config{})
	require.NoError(t, err)
	require.NoError(t, db.AutoMigrate(&types.MemorySubject{}))
	for i := range 50 {
		require.NoError(t, db.Create(&types.MemorySubject{ID: fmt.Sprint(i), TenantID: 1, SubjectID: fmt.Sprintf("web_user:%d", i), PendingSessions: types.MemoryPendingSessions{"s"}}).Error)
	}
	require.NoError(t, db.Create(&types.MemorySubject{ID: "idle", TenantID: 1, SubjectID: "web_user:idle"}).Error)
	repo := NewMemoryRepository(db)
	due, err := repo.ListDueExtractions(context.Background(), 1000)
	require.NoError(t, err)
	require.Len(t, due, 32)
	for _, subject := range due {
		require.NotEmpty(t, subject.PendingSessions)
	}
}

func TestSQLiteMemoryRecoveryMigrationRoundTrip(t *testing.T) {
	db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
	require.NoError(t, err)
	require.NoError(t, db.Exec("CREATE TABLE memory_subjects (id TEXT PRIMARY KEY, pending_sessions JSON)").Error)
	for _, direction := range []string{"up", "down", "up"} {
		migration, err := os.ReadFile("../../../migrations/sqlite/000217_memory_extraction_recovery." + direction + ".sql")
		require.NoError(t, err)
		require.NoError(t, db.Exec(string(migration)).Error)
	}
	require.True(t, db.Migrator().HasColumn("memory_subjects", "extract_failures"))
	require.True(t, db.Migrator().HasColumn("memory_subjects", "session_cursors"))
}

func TestMemoryRecoveryItemWritesRequireLiveReservation(t *testing.T) {
	for _, change := range []string{"superseded", "expired", "disabled", "deleted", "not started"} {
		t.Run(change, func(t *testing.T) {
			db, err := gorm.Open(sqlite.Open("file:"+t.Name()+"?mode=memory&cache=shared"), &gorm.Config{})
			require.NoError(t, err)
			require.NoError(t, db.AutoMigrate(&types.MemorySubject{}, &types.MemoryItem{}))
			repo := NewMemoryRepository(db)
			ctx := context.Background()
			scope := interfaces.MemoryScope{TenantID: 1, SubjectID: "web_user:alice"}
			_, err = repo.EnsureSubject(ctx, scope)
			require.NoError(t, err)
			_, _, err = repo.EnqueuePendingSession(ctx, scope, "s", time.Minute)
			require.NoError(t, err)
			payload, err := repo.ReserveExtraction(ctx, scope, types.MemoryExtractPayload{}, 0, time.Minute)
			require.NoError(t, err)
			require.NotNil(t, payload)
			started, err := repo.StartExtraction(ctx, scope, payload.AttemptID, time.Minute)
			require.NoError(t, err)
			require.True(t, started)
			workerCtx := types.WithMemoryExtractionAttempt(ctx, payload.AttemptID)
			original := &types.MemoryItem{TenantID: 1, SubjectID: scope.SubjectID, Kind: types.MemoryKindFact, Content: "original"}
			require.NoError(t, repo.CreateItem(workerCtx, original))
			switch change {
			case "superseded":
				err = db.Model(&types.MemorySubject{}).Where("tenant_id = ?", 1).Update("extract_attempt_id", "new-attempt").Error
			case "expired":
				err = db.Model(&types.MemorySubject{}).Where("tenant_id = ?", 1).Update("extract_lease_until", time.Now().Add(-time.Second)).Error
			case "disabled":
				err = repo.UpdateSubjectEnabled(ctx, scope, false)
			case "deleted":
				err = db.Where("tenant_id = ?", 1).Delete(&types.MemorySubject{}).Error
			case "not started":
				err = db.Model(&types.MemorySubject{}).Where("tenant_id = ?", 1).Update("extract_started_at", nil).Error
			}
			require.NoError(t, err)
			require.Error(t, repo.CreateItem(workerCtx, &types.MemoryItem{TenantID: 1, SubjectID: scope.SubjectID, Content: "stale"}))
			require.Error(t, repo.SupersedeItem(workerCtx, scope, original.ID, ""))
			_, err = repo.EnsureSubject(workerCtx, scope)
			require.Error(t, err)
			var count int64
			require.NoError(t, db.Model(&types.MemoryItem{}).Count(&count).Error)
			require.EqualValues(t, 1, count)
			unchanged, err := repo.GetItem(ctx, scope, original.ID)
			require.NoError(t, err)
			require.Equal(t, types.MemoryStatusActive, unchanged.Status)
			if change == "deleted" {
				subject, err := repo.GetSubject(ctx, scope)
				require.NoError(t, err)
				require.Nil(t, subject)
			}
		})
	}
}