xinmi/internal/workbench/memory_legacy_postgres_test.go

package workbench

import (
	"os"
	"strings"
	"sync"
	"testing"

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

func TestPostgresLegacyMemoryMigrationProvenance(t *testing.T) {
	db := postgresLegacyMemoryDB(t)
	if err := db.Exec(`CREATE TABLE workbench_pending_memories (
        id varchar(36) PRIMARY KEY, tenant_id bigint NOT NULL, owner_id varchar(36) NOT NULL,
        agent_id varchar(64) NOT NULL, content text NOT NULL, source_ref text NOT NULL,
        status varchar(32) NOT NULL, reviewed_by varchar(36) NOT NULL DEFAULT '',
        created_at timestamptz NOT NULL DEFAULT now(), updated_at timestamptz NOT NULL DEFAULT now());
        CREATE TABLE workbench_history (id varchar(36) PRIMARY KEY, tenant_id bigint, source_kind text, source_id text);
        CREATE TABLE memory_subjects (tenant_id bigint, subject_id text, enabled boolean NOT NULL DEFAULT true,
            pending_sessions jsonb NOT NULL DEFAULT '[]');
        CREATE TABLE yawan_legacy_operations_mappings (history_id varchar(36), target_tenant_id bigint,
            source_kind text, source_id bigint, target_table text, target_key text, native_snapshot jsonb);
        INSERT INTO workbench_pending_memories(id,tenant_id,owner_id,agent_id,content,source_ref,status)
        SELECT n::text,1,'alice','author','original advice','original source','accepted' FROM generate_series(1,6) n;
        INSERT INTO workbench_history SELECT n::text,1,'pendingmemory',n::text FROM generate_series(1,5) n;
        INSERT INTO yawan_legacy_operations_mappings
        SELECT n::text,1,'pendingmemory',n,'workbench_pending_memories',n::text,
        '{"owner_id":"alice","agent_id":"author","status":"accepted"}'::jsonb FROM generate_series(1,6) n;
        UPDATE yawan_legacy_operations_mappings SET native_snapshot=jsonb_set(native_snapshot,'{owner_id}','"bob"') WHERE source_id=2;
        UPDATE yawan_legacy_operations_mappings SET target_tenant_id=2 WHERE source_id=3;
        UPDATE yawan_legacy_operations_mappings SET native_snapshot=jsonb_set(native_snapshot,'{status}','"pending"') WHERE source_id=4;
        UPDATE workbench_pending_memories SET status='rejected' WHERE id='5';`).Error; err != nil {
		t.Fatal(err)
	}
	raw, err := os.ReadFile("../../migrations/versioned/000217_memory_recovery_reconfirmation.up.sql")
	if err != nil {
		t.Fatal(err)
	}
	for i := 0; i < 2; i++ {
		if err := db.Exec(string(raw)).Error; err != nil {
			t.Fatal(err)
		}
	}
	var rows []PendingMemory
	if err := db.Order("id").Find(&rows).Error; err != nil {
		t.Fatal(err)
	}
	if len(rows) != 6 {
		t.Fatalf("history disappeared: %d", len(rows))
	}
	for _, row := range rows {
		if (row.LegacyHistoryID != "") != (row.ID == "1") || row.NativeMemoryID != "" || row.ReconfirmedAt != nil || row.Content != "original advice" || row.SourceRef != "original source" {
			t.Fatalf("migration changed data or admitted wrong provenance: %+v", row)
		}
	}
	if err := db.Exec("DROP TABLE yawan_legacy_operations_mappings").Error; err != nil {
		t.Fatal(err)
	}
	if err := db.Exec(string(raw)).Error; err != nil {
		t.Fatalf("native-only installation must migrate: %v", err)
	}
}

func TestPostgresLegacyMemoryReconfirmationConcurrent(t *testing.T) {
	db := postgresLegacyMemoryDB(t)
	if err := db.AutoMigrate(&PendingMemory{}); err != nil {
		t.Fatal(err)
	}
	s, ctx, row, calls := legacyMemoryHarnessWithService(t, NewService(db))
	var wg sync.WaitGroup
	codes := make(chan int, 2)
	for i := 0; i < 2; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			w := legacyMemoryRequest(s, ctx, row.ID, memoryReviewRequest{Status: "accepted", Content: "stable preference", ReconfirmLegacy: true, ConfirmPersonalScope: true})
			codes <- w.Code
		}()
	}
	wg.Wait()
	close(codes)
	for code := range codes {
		if code != 200 {
			t.Fatalf("concurrent confirmation failed: %d", code)
		}
	}
	if *calls != 1 {
		t.Fatalf("concurrent retries wrote %d native memories", *calls)
	}
}

func postgresLegacyMemoryDB(t *testing.T) *gorm.DB {
	t.Helper()
	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)
	}
	schema := "wb_legacy_memory_" + strings.ReplaceAll(uuid.NewString(), "-", "")
	if err := db.Exec("CREATE SCHEMA " + schema).Error; err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() {
		if err := db.Exec("DROP SCHEMA " + schema + " CASCADE").Error; err != nil {
			t.Error(err)
		}
		_ = sqlDB.Close()
	})
	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)
	}
	t.Cleanup(func() { _ = testSQL.Close() })
	return testDB
}