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 }