xinmi/internal/workbench/memory_review_test.go

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)
			}
		})
	}
}