package workbench import ( "bytes" "context" "errors" "net/http" "net/http/httptest" "strings" "testing" "github.com/Tencent/WeKnora/internal/types" "github.com/gin-gonic/gin" "gorm.io/gorm" ) func TestMemoryApprovalIsAtomicAndOwnerScoped(t *testing.T) { for _, mode := range []string{"accepted", "rejected", "write_failure", "status_failure", "nil_result", "pending_result", "other_owner", "other_tenant"} { t.Run(mode, func(t *testing.T) { s := sqliteWorkbench(t) if err := s.db.AutoMigrate(&types.MemoryItem{}); err != nil { t.Fatal(err) } row := PendingMemory{ID: "pending", TenantID: 1, OwnerID: "alice", Content: "每次直播先介绍主题", Status: "pending"} if err := s.db.Create(&row).Error; err != nil { t.Fatal(err) } s.SetTransactionalMemoryWriter(func(ctx context.Context, tx *gorm.DB, item types.MemoryItem) (*types.MemoryItem, error) { item.ID, item.TenantID, item.SubjectID = "native", 1, "alice" item.Status = types.MemoryStatusActive if err := tx.Create(&item).Error; err != nil { return nil, err } if mode == "write_failure" { return nil, errors.New("write failed") } if mode == "nil_result" { return nil, nil } if mode == "pending_result" { item.Status = types.MemoryStatusPending } return &item, nil }) if mode == "status_failure" { if err := s.db.Exec("CREATE TRIGGER reject_approval BEFORE UPDATE ON workbench_pending_memories BEGIN SELECT RAISE(ABORT, 'status failed'); END").Error; err != nil { t.Fatal(err) } } ctx := workbenchContext(1, "alice", types.TenantRoleContributor) if mode == "other_owner" { ctx = workbenchContext(1, "bob", types.TenantRoleAdmin) } if mode == "other_tenant" { ctx = workbenchContext(2, "alice", types.TenantRoleAdmin) } status := "accepted" if mode == "rejected" { status = "rejected" } w := httptest.NewRecorder() c, _ := gin.CreateTestContext(w) c.Params = gin.Params{{Key: "id", Value: row.ID}} c.Request = httptest.NewRequest(http.MethodPost, "/", bytes.NewBufferString(`{"status":"`+status+`"}`)).WithContext(ctx) c.Request.Header.Set("Content-Type", "application/json") NewHandler(s).ReviewMemory(c) var count int64 if err := s.db.Model(&types.MemoryItem{}).Count(&count).Error; err != nil { t.Fatal(err) } if err := s.db.First(&row, "id = ?", row.ID).Error; err != nil { t.Fatal(err) } if mode == "accepted" || mode == "rejected" { if w.Code != 200 || row.Status != status { t.Fatalf("approval failed: %d %s", w.Code, w.Body.String()) } if (mode == "accepted" && count != 1) || (mode == "rejected" && count != 0) { t.Fatalf("unexpected memories: %d", count) } } else if w.Code < 400 || row.Status != "pending" || count != 0 { t.Fatalf("failed approval leaked state: http=%d pending=%s memories=%d", w.Code, row.Status, count) } }) } } func TestMemoryApprovalLengthAndRetry(t *testing.T) { for _, mode := range []string{"long_accept", "long_reject", "retry", "sanitized"} { t.Run(mode, func(t *testing.T) { s := sqliteWorkbench(t) if err := s.db.AutoMigrate(&types.MemoryItem{}); err != nil { t.Fatal(err) } row := PendingMemory{ID: "pending", TenantID: 1, OwnerID: "alice", Content: "每次直播\n先介绍主题", Status: "pending"} if strings.HasPrefix(mode, "long_") { row.Content = strings.Repeat("字", types.MemoryContentMaxRunes+1) } if err := s.db.Create(&row).Error; err != nil { t.Fatal(err) } calls := 0 s.SetTransactionalMemoryWriter(func(ctx context.Context, tx *gorm.DB, item types.MemoryItem) (*types.MemoryItem, error) { calls++ if mode == "retry" && calls == 1 { return nil, errors.New("temporary write failure") } item.ID, item.TenantID, item.SubjectID = "native", 1, "alice" item.Status = types.MemoryStatusActive item.Content = types.SanitizeMemoryContent(item.Content) return &item, tx.Create(&item).Error }) request := func(status string) *httptest.ResponseRecorder { w := httptest.NewRecorder() c, _ := gin.CreateTestContext(w) c.Params = gin.Params{{Key: "id", Value: row.ID}} ctx := workbenchContext(1, "alice", types.TenantRoleContributor) c.Request = httptest.NewRequest(http.MethodPost, "/", bytes.NewBufferString(`{"status":"`+status+`"}`)).WithContext(ctx) c.Request.Header.Set("Content-Type", "application/json") NewHandler(s).ReviewMemory(c) return w } status := "accepted" if mode == "long_reject" { status = "rejected" } w := request(status) if mode == "retry" { if w.Code < 400 { t.Fatal("failed write reported success") } if err := s.db.First(&row, "id = ?", row.ID).Error; err != nil || row.Status != "pending" { t.Fatalf("retry lost pending row: %+v %v", row, err) } w = request(status) } if err := s.db.First(&row, "id = ?", row.ID).Error; err != nil { t.Fatal(err) } var count int64 if err := s.db.Model(&types.MemoryItem{}).Count(&count).Error; err != nil { t.Fatal(err) } if mode == "long_accept" { if w.Code != http.StatusBadRequest || !strings.Contains(w.Body.String(), "300") || row.Status != "pending" || calls != 0 || count != 0 || len([]rune(row.Content)) != types.MemoryContentMaxRunes+1 { t.Fatalf("long approval was truncated or lost: response=%d %s row=%+v", w.Code, w.Body.String(), row) } return } if w.Code != http.StatusOK || row.Status != status { t.Fatalf("review failed: %d %s row=%+v", w.Code, w.Body.String(), row) } if mode == "long_reject" { if calls != 0 || count != 0 { t.Fatal("rejection created memory") } } else if count != 1 || row.Content != "每次直播 先介绍主题" { t.Fatalf("approved history differs from stored memory: count=%d row=%+v", count, row) } }) } }