xinmi/internal/application/service/knowledge_batch_access_test.go

package service

import (
	"context"
	"testing"

	"github.com/Tencent/WeKnora/internal/accesspolicy"
	"github.com/Tencent/WeKnora/internal/types"
	"github.com/Tencent/WeKnora/internal/types/interfaces"
)

type batchAccessKBStub struct {
	interfaces.KnowledgeBaseService
	calls map[string]int
}

func (s *batchAccessKBStub) GetKnowledgeBaseByIDOnly(_ context.Context, id string) (*types.KnowledgeBase, error) {
	s.calls[id]++
	return &types.KnowledgeBase{ID: id, TenantID: 7, CreatorID: "owner"}, nil
}

func TestFilterAccessibleKnowledgeHonorsGrantAndRevocation(t *testing.T) {
	kb := &batchAccessKBStub{calls: map[string]int{}}
	s := &knowledgeService{kbService: kb}
	ctx := context.WithValue(context.Background(), types.TenantIDContextKey, uint64(7))
	ctx = context.WithValue(ctx, types.UserIDContextKey, "reader")
	policy := &accesspolicy.Snapshot{Restricted: true, Permissions: map[string]bool{"knowledge.read": true}, Publications: map[string]bool{"knowledge:private": false, "knowledge:public": true}, Grants: map[string]bool{"knowledge:private": true}}
	ctx = accesspolicy.WithSnapshot(ctx, policy)
	rows := []*types.Knowledge{{ID: "one", KnowledgeBaseID: "private"}, {ID: "two", KnowledgeBaseID: "private"}, {ID: "public", KnowledgeBaseID: "public"}, {ID: "unpublished", KnowledgeBaseID: "unpublished"}, nil}
	if got := s.filterAccessibleKnowledge(ctx, rows); len(got) != 3 || kb.calls["private"] != 1 {
		t.Fatalf("grant should allow documents and check each KB once: %d rows, %v calls", len(got), kb.calls)
	}
	delete(policy.Grants, "knowledge:private")
	if got := s.filterAccessibleKnowledge(ctx, rows); len(got) != 1 || got[0].ID != "public" {
		t.Fatalf("revoked documents remain in batch: %v", got)
	}
	delete(policy.Permissions, "knowledge.read")
	if got := s.filterAccessibleKnowledge(ctx, rows); len(got) != 0 {
		t.Fatalf("module permission must restrict all documents: %v", got)
	}
}