xinmi/internal/workbench/content_memory_postgres_test.go

package workbench

import (
	"os"
	"strings"
	"testing"
	"time"

	"github.com/google/uuid"
	"gorm.io/driver/postgres"
	"gorm.io/gorm"
	"gorm.io/gorm/logger"
)

// PostgreSQL distinguishes the conflict target from EXCLUDED; SQLite does not
// reject an unqualified status predicate, so this must execute on PostgreSQL.
func TestPostgresContentFeedbackAcceptanceUpsert(t *testing.T) {
	dsn := os.Getenv("WORKBENCH_TEST_POSTGRES_DSN")
	if dsn == "" {
		t.Skip("WORKBENCH_TEST_POSTGRES_DSN is not set")
	}
	db, err := gorm.Open(postgres.Open(dsn), &gorm.Config{Logger: logger.Default.LogMode(logger.Silent)})
	if err != nil {
		t.Fatal(err)
	}
	sqlDB, err := db.DB()
	if err != nil {
		t.Fatal(err)
	}
	defer sqlDB.Close()
	schema := "wb_memory_test_" + strings.ReplaceAll(uuid.NewString(), "-", "")
	if err := db.Exec("CREATE SCHEMA " + schema).Error; err != nil {
		t.Fatal(err)
	}
	defer func() {
		if err := db.Exec("DROP SCHEMA " + schema + " CASCADE").Error; err != nil {
			t.Error(err)
		}
	}()
	testDSN := dsn + " search_path=" + schema
	if strings.Contains(dsn, "://") {
		separator := "?"
		if strings.Contains(dsn, "?") {
			separator = "&"
		}
		testDSN = dsn + separator + "search_path=" + schema
	}
	testDB, err := gorm.Open(postgres.Open(testDSN), &gorm.Config{Logger: logger.Default.LogMode(logger.Silent)})
	if err != nil {
		t.Fatal(err)
	}
	testSQL, err := testDB.DB()
	if err != nil {
		t.Fatal(err)
	}
	defer testSQL.Close()
	if err := testDB.AutoMigrate(&Resource{}, &Job{}, &Revision{}, &Event{}, &Consumption{}, &PendingMemory{}); err != nil {
		t.Fatal(err)
	}
	s, ctx, resource, job := feedbackHarnessWithService(t, NewService(testDB))
	accept := func() PendingMemory {
		t.Helper()
		version := resource.Version
		if err := s.executeTask(ctx, job, resource, ActionRequest{SegmentID: "segment"}); err != nil {
			t.Fatal(err)
		}
		if err := s.save(ctx, resource, version); err != nil {
			t.Fatalf("acceptance must commit with its pending feedback: %v", err)
		}
		var persisted Resource
		if err := testDB.First(&persisted, "id = ?", resource.ID).Error; err != nil {
			t.Fatal(err)
		}
		data, err := decode[TaskData](persisted.Data)
		if err != nil || persisted.Version != version+1 || data.Segments[0].Status != "accepted" {
			t.Fatalf("acceptance was not persisted atomically: %+v %v", persisted, err)
		}
		var rows []PendingMemory
		if err := testDB.Where("tenant_id = ? AND owner_id = ?", resource.TenantID, resource.OwnerID).Find(&rows).Error; err != nil {
			t.Fatal(err)
		}
		if len(rows) != 1 || rows[0].Content != "开场先说结论，再给案例" {
			t.Fatalf("feedback lost or duplicated: %+v", rows)
		}
		return rows[0]
	}
	first := accept()
	if first.Status != "pending" {
		t.Fatalf("initial acceptance did not create pending feedback: %+v", first)
	}
	if repeated := accept(); repeated.ID != first.ID || !repeated.UpdatedAt.Equal(first.UpdatedAt) {
		t.Fatalf("repeat acceptance mutated an existing pending candidate: %+v", repeated)
	}
	for _, state := range []string{"invalidated", "accepted", "rejected"} {
		stamp := time.Date(2020, 1, 2, 3, 4, 5, 0, time.UTC)
		if err := testDB.Model(&PendingMemory{}).Where("id = ?", first.ID).
			Updates(map[string]any{"status": state, "updated_at": stamp}).Error; err != nil {
			t.Fatal(err)
		}
		row := accept()
		if state == "invalidated" {
			if row.Status != "pending" || !row.UpdatedAt.After(stamp) {
				t.Fatalf("reacceptance did not restore invalidated feedback: %+v", row)
			}
		} else if row.Status != state || !row.UpdatedAt.Equal(stamp) {
			t.Fatalf("reacceptance changed explicit %s review: %+v", state, row)
		}
	}
}