xinmi/internal/middleware/knowledge_export_test.go

package middleware

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

	"github.com/Tencent/WeKnora/internal/accesspolicy"
	"github.com/Tencent/WeKnora/internal/types"
	"github.com/gin-gonic/gin"
	"github.com/stretchr/testify/require"
)

func TestRequireKnowledgeExportFailsClosedWithoutResolvedKnowledgeBase(t *testing.T) {
	gin.SetMode(gin.TestMode)
	for _, tc := range []struct {
		name   string
		access any
	}{
		{name: "absent"},
		{name: "nil access", access: (*KBAccess)(nil)},
		{name: "missing KB", access: &KBAccess{}},
		{name: "invalid type", access: "kb-id"},
	} {
		t.Run(tc.name, func(t *testing.T) {
			r := gin.New()
			r.Use(ErrorHandler())
			r.Use(func(c *gin.Context) {
				ctx := accesspolicy.WithSnapshot(c.Request.Context(), &accesspolicy.Snapshot{
					Publications: map[string]bool{"knowledge:kb-id": true},
				})
				c.Request = c.Request.WithContext(ctx)
				if tc.access != nil {
					c.Set(KBAccessContextKey, tc.access)
				}
				c.Next()
			})
			r.GET("/export", RequireKnowledgeExport(), func(c *gin.Context) {
				t.Fatal("export handler ran without an authorized KB")
			})
			rec := httptest.NewRecorder()
			r.ServeHTTP(rec, httptest.NewRequest(http.MethodGet, "/export", nil))
			require.Equal(t, http.StatusServiceUnavailable, rec.Code)
		})
	}
}

func TestRequireKnowledgeExportPreservesSystemAdministrator(t *testing.T) {
	gin.SetMode(gin.TestMode)
	r := gin.New()
	r.Use(func(c *gin.Context) {
		ctx := context.WithValue(c.Request.Context(), types.SystemAdminContextKey, true)
		ctx = context.WithValue(ctx, types.TenantRoleContextKey, types.TenantRoleViewer)
		ctx = accesspolicy.WithSnapshot(ctx, &accesspolicy.Snapshot{Publications: map[string]bool{"knowledge:kb-id": true}})
		c.Request = c.Request.WithContext(ctx)
		c.Set(KBAccessContextKey, &KBAccess{KnowledgeBase: &types.KnowledgeBase{ID: "kb-id"}})
		c.Next()
	})
	r.GET("/export", RequireKnowledgeExport(), func(c *gin.Context) { c.Status(http.StatusNoContent) })
	rec := httptest.NewRecorder()
	r.ServeHTTP(rec, httptest.NewRequest(http.MethodGet, "/export", nil))
	require.Equal(t, http.StatusNoContent, rec.Code)
}