76 lines
3.2 KiB
Go
76 lines
3.2 KiB
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"}))
|
||
|
|
}
|