xinmi/internal/handler/knowledge_batch_ownership_test.go

package handler

import (
	"context"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

	"github.com/Tencent/WeKnora/internal/config"
	"github.com/Tencent/WeKnora/internal/middleware"
	"github.com/Tencent/WeKnora/internal/types"
	"github.com/Tencent/WeKnora/internal/types/interfaces"
	"github.com/gin-gonic/gin"
	"github.com/hibiken/asynq"
	"github.com/stretchr/testify/require"
)

type batchOwnershipKBService struct {
	interfaces.KnowledgeBaseService
}

func (*batchOwnershipKBService) GetKnowledgeBaseByID(context.Context, string) (*types.KnowledgeBase, error) {
	return &types.KnowledgeBase{ID: "kb-owner", TenantID: 1, CreatorID: "owner"}, nil
}

type batchOwnershipKnowledgeService struct {
	interfaces.KnowledgeService
	updated bool
	fetched bool
}

func (*batchOwnershipKnowledgeService) GetKnowledgeByIDOnly(context.Context, string) (*types.Knowledge, error) {
	return &types.Knowledge{ID: "doc-owner", TenantID: 1, KnowledgeBaseID: "kb-owner"}, nil
}

func (s *batchOwnershipKnowledgeService) UpdateKnowledgeTagBatch(_ context.Context, kbID string, _ map[string][]string) error {
	s.updated = kbID == "kb-owner"
	return nil
}

func (s *batchOwnershipKnowledgeService) GetKnowledgeBatch(context.Context, uint64, []string) ([]*types.Knowledge, error) {
	s.fetched = true
	return []*types.Knowledge{{ID: "doc-owner", TenantID: 1, KnowledgeBaseID: "kb-owner"}}, nil
}

type batchOwnershipQueue struct{ enqueued bool }

func (q *batchOwnershipQueue) Enqueue(*asynq.Task, ...asynq.Option) (*asynq.TaskInfo, error) {
	q.enqueued = true
	return &asynq.TaskInfo{ID: "reparse-test"}, nil
}

func TestKnowledgeBatchMutationsRequireKBOwnerOrAdmin(t *testing.T) {
	gin.SetMode(gin.TestMode)
	for _, principal := range []struct {
		name string
		user string
		role types.TenantRole
		want int
	}{
		{"other contributor", "other", types.TenantRoleContributor, http.StatusForbidden},
		{"other viewer", "other", types.TenantRoleViewer, http.StatusForbidden},
		{"KB creator", "owner", types.TenantRoleContributor, http.StatusOK},
		{"tenant admin", "other", types.TenantRoleAdmin, http.StatusOK},
	} {
		for _, operation := range []struct {
			name string
			body string
		}{
			{"tags with KB", `{"kb_id":"kb-owner","updates":{"doc-owner":["tag-1"]}}`},
			{"tags inferred KB", `{"updates":{"doc-owner":["tag-1"]}}`},
			{"batch reparse", `{"kb_id":"kb-owner","ids":["doc-owner"]}`},
		} {
			t.Run(principal.name+"/"+operation.name, func(t *testing.T) {
				enabled := true
				cfg := &config.Config{Tenant: &config.TenantConfig{EnableRBAC: &enabled}}
				kg := &batchOwnershipKnowledgeService{}
				queue := &batchOwnershipQueue{}
				h := &KnowledgeHandler{cfg: cfg, kbService: &batchOwnershipKBService{}, kgService: kg, asynqClient: queue}
				r := gin.New()
				r.Use(middleware.ErrorHandler())
				r.Use(func(c *gin.Context) {
					ctx := context.WithValue(c.Request.Context(), types.TenantIDContextKey, uint64(1))
					ctx = context.WithValue(ctx, types.UserIDContextKey, principal.user)
					ctx = context.WithValue(ctx, types.TenantRoleContextKey, principal.role)
					c.Request = c.Request.WithContext(ctx)
					c.Set(types.TenantIDContextKey.String(), uint64(1))
					c.Set(types.UserIDContextKey.String(), principal.user)
					c.Next()
				})
				if operation.name == "batch reparse" {
					r.POST("/batch", h.BatchReparseKnowledge)
				} else {
					r.POST("/batch", h.UpdateKnowledgeTagBatch)
				}
				req := httptest.NewRequest(http.MethodPost, "/batch", strings.NewReader(operation.body))
				req.Header.Set("Content-Type", "application/json")
				w := httptest.NewRecorder()
				r.ServeHTTP(w, req)
				require.Equal(t, principal.want, w.Code, "body=%s", w.Body.String())
				allowed := principal.want == http.StatusOK
				if operation.name == "batch reparse" {
					require.Equal(t, allowed, kg.fetched, "reject before reading batch")
					require.Equal(t, allowed, queue.enqueued, "reject before scheduling reparse")
				} else {
					require.Equal(t, allowed, kg.updated, "reject before changing tags")
				}
			})
		}
	}
}