xinmi/internal/application/service/custom_agent_suggestion_access_test.go

package service

import (
	"context"
	"errors"
	"testing"

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

type suggestionAccessKnowledgeRepo struct {
	interfaces.KnowledgeRepository
	documents map[string]*types.Knowledge
}

func (r suggestionAccessKnowledgeRepo) GetKnowledgeByIDOnly(_ context.Context, id string) (*types.Knowledge, error) {
	return r.documents[id], nil
}

type suggestionAccessKBService struct {
	interfaces.KnowledgeBaseService
	bases map[string]*types.KnowledgeBase
	err   error
}

func (s suggestionAccessKBService) GetKnowledgeBasesByIDsOnly(ctx context.Context, ids []string) ([]*types.KnowledgeBase, error) {
	if s.err != nil {
		return nil, s.err
	}
	var result []*types.KnowledgeBase
	for _, id := range ids {
		if kb := s.bases[id]; kb != nil {
			result = append(result, kb)
		}
	}
	return filterReadableKnowledgeBases(ctx, result), nil
}

func TestSuggestionDocumentIDsRespectCurrentKBAccess(t *testing.T) {
	ctx := context.WithValue(context.Background(), types.TenantIDContextKey, uint64(1))
	ctx = context.WithValue(ctx, types.UserIDContextKey, "caller")
	ctx = context.WithValue(ctx, types.TenantRoleContextKey, types.TenantRoleContributor)
	svc := &customAgentService{
		knowledgeRepo: suggestionAccessKnowledgeRepo{documents: map[string]*types.Knowledge{
			"own-doc":     {ID: "own-doc", TenantID: 1, KnowledgeBaseID: "own"},
			"private-doc": {ID: "private-doc", TenantID: 1, KnowledgeBaseID: "private"},
			"shared-doc":  {ID: "shared-doc", TenantID: 2, KnowledgeBaseID: "shared"},
			"foreign-doc": {ID: "foreign-doc", TenantID: 2, KnowledgeBaseID: "foreign"},
		}},
		kbService: suggestionAccessKBService{bases: map[string]*types.KnowledgeBase{
			"own":     {ID: "own", TenantID: 1, CreatorID: "caller"},
			"private": {ID: "private", TenantID: 1, CreatorID: "other"},
			"shared":  {ID: "shared", TenantID: 2},
			"foreign": {ID: "foreign", TenantID: 2},
		}},
		kbShareService: &suggestionKBShareService{allowed: map[string]bool{"shared": true}},
	}
	ids := []string{"own-doc", "private-doc", "shared-doc", "foreign-doc", "missing", "own-doc"}
	private := accesspolicy.WithSnapshot(ctx, &accesspolicy.Snapshot{})
	require.Equal(t, map[uint64][]string{1: {"own-doc"}, 2: {"shared-doc"}}, svc.groupSuggestionKnowledgeByTenant(private, 1, ids))
	granted := accesspolicy.WithSnapshot(ctx, &accesspolicy.Snapshot{Grants: map[string]bool{"knowledge:private": true}})
	require.Equal(t, []string{"own-doc", "private-doc"}, svc.groupSuggestionKnowledgeByTenant(granted, 1, ids)[1])
	revoked := accesspolicy.WithSnapshot(ctx, &accesspolicy.Snapshot{Publications: map[string]bool{"knowledge:private": false}})
	require.Empty(t, svc.groupSuggestionKnowledgeByTenant(revoked, 1, []string{"private-doc"}))
	denied := accesspolicy.WithSnapshot(ctx, &accesspolicy.Snapshot{Restricted: true})
	require.Empty(t, svc.groupSuggestionKnowledgeByTenant(denied, 1, ids))
}

func TestSuggestionKBResolutionFailsClosed(t *testing.T) {
	svc := &customAgentService{kbService: suggestionAccessKBService{err: errors.New("database unavailable")}}
	require.Empty(t, svc.groupKBIDsByEffectiveTenant(context.Background(), 1, []string{"private"}))
}