79 lines
2.7 KiB
Go
79 lines
2.7 KiB
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"
|
||
|
|
"github.com/stretchr/testify/require"
|
||
|
|
)
|
||
|
|
|
||
|
|
type executionShareStub struct {
|
||
|
|
interfaces.AgentShareService
|
||
|
|
agent *types.CustomAgent
|
||
|
|
err error
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *executionShareStub) GetSharedAgentForTenant(_ context.Context, tenant uint64, _ types.TenantRole, id string, source ...uint64) (*types.CustomAgent, error) {
|
||
|
|
if tenant != 2 || id != "root" || len(source) != 1 || source[0] != 1 {
|
||
|
|
return nil, ErrAgentSharePermission
|
||
|
|
}
|
||
|
|
return s.agent, s.err
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestSessionAgentExecutionLoadsCurrentBindings(t *testing.T) {
|
||
|
|
agents, repo, kbs, ctx := agentExecutionHarness(t)
|
||
|
|
old := *repo.rows["root"]
|
||
|
|
repo.rows["root"].Config.KnowledgeBases = []string{"nested"}
|
||
|
|
s := &sessionService{customAgentService: agents}
|
||
|
|
req := &types.QARequest{Session: &types.Session{TenantID: 1}, CustomAgent: &old}
|
||
|
|
run, err := s.prepareQAAgentExecution(ctx, req)
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.Equal(t, []string{"nested"}, req.CustomAgent.Config.KnowledgeBases)
|
||
|
|
_, err = kbs.GetKnowledgeBaseByID(run, "nested")
|
||
|
|
require.NoError(t, err)
|
||
|
|
_, err = kbs.GetKnowledgeBaseByID(run, "owned")
|
||
|
|
require.Error(t, err, "old request bindings must not survive a live edit")
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestSessionAgentExecutionRechecksOrganizationShare(t *testing.T) {
|
||
|
|
agents, repo, kbs, ctx := agentExecutionHarness(t)
|
||
|
|
accesspolicy.FromContext(ctx).Grants = nil
|
||
|
|
share := &executionShareStub{agent: repo.rows["root"]}
|
||
|
|
s := &sessionService{customAgentService: agents, agentShareService: share}
|
||
|
|
req := &types.QARequest{Session: &types.Session{TenantID: 2}, CustomAgent: repo.rows["root"], SharedAgentReadOnly: true}
|
||
|
|
run, err := s.prepareQAAgentExecution(ctx, req)
|
||
|
|
require.NoError(t, err)
|
||
|
|
_, err = kbs.GetKnowledgeBaseByID(run, "owned")
|
||
|
|
require.NoError(t, err)
|
||
|
|
_, err = kbs.GetKnowledgeBaseByID(ctx, "owned")
|
||
|
|
require.Error(t, err)
|
||
|
|
share.err = ErrAgentSharePermission
|
||
|
|
_, err = s.prepareQAAgentExecution(ctx, req)
|
||
|
|
require.Error(t, err)
|
||
|
|
share.err = nil
|
||
|
|
repo.rows["root"].Enabled = false
|
||
|
|
_, err = s.prepareQAAgentExecution(ctx, req)
|
||
|
|
require.Error(t, err)
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestSessionQABlocksRevokedAgentBeforeModelInvocation(t *testing.T) {
|
||
|
|
for _, mode := range []string{"agent", "knowledge"} {
|
||
|
|
t.Run(mode, func(t *testing.T) {
|
||
|
|
agents, repo, _, ctx := agentExecutionHarness(t)
|
||
|
|
accesspolicy.FromContext(ctx).Grants = nil
|
||
|
|
s := &sessionService{customAgentService: agents}
|
||
|
|
req := &types.QARequest{Session: &types.Session{ID: "test", TenantID: 1}, CustomAgent: repo.rows["root"]}
|
||
|
|
var err error
|
||
|
|
if mode == "agent" {
|
||
|
|
err = s.AgentQA(ctx, req, nil)
|
||
|
|
} else {
|
||
|
|
err = s.KnowledgeQA(ctx, req, nil)
|
||
|
|
}
|
||
|
|
require.Error(t, err)
|
||
|
|
})
|
||
|
|
}
|
||
|
|
}
|