package service import ( "context" "encoding/json" "errors" "strings" "testing" "time" "github.com/Tencent/WeKnora/internal/accesspolicy" "github.com/Tencent/WeKnora/internal/models/chat" "github.com/Tencent/WeKnora/internal/types" "github.com/Tencent/WeKnora/internal/types/interfaces" "github.com/stretchr/testify/require" ) type digitalHumanTrialModelService struct { interfaces.ModelService model chat.Chat status types.ModelStatus } func (s *digitalHumanTrialModelService) GetModelByID(context.Context, string) (*types.Model, error) { return &types.Model{ID: "configured-model", Type: types.ModelTypeKnowledgeQA, Status: s.status}, nil } func (s *digitalHumanTrialModelService) GetChatModel(context.Context, string) (chat.Chat, error) { return s.model, nil } type digitalHumanTrialChat struct { messages []chat.Message purpose string owner string tenant uint64 chunks []types.StreamResponse err error block bool cancelled chan struct{} } func (m *digitalHumanTrialChat) GetModelID() string { return "configured-model" } func (m *digitalHumanTrialChat) GetModelName() string { return "trial-fixture" } func (m *digitalHumanTrialChat) Chat(context.Context, []chat.Message, *chat.ChatOptions) (*types.ChatResponse, error) { return nil, errors.New("trial must use the metered stream") } func (m *digitalHumanTrialChat) ChatStream(ctx context.Context, messages []chat.Message, options *chat.ChatOptions) (<-chan types.StreamResponse, error) { m.messages = messages m.purpose, _ = types.LLMCallMetadataFromContext(ctx) m.owner, _ = types.UserIDFromContext(ctx) m.tenant, _ = types.TenantIDFromContext(ctx) if m.err != nil { return nil, m.err } if m.block { <-ctx.Done() close(m.cancelled) return nil, ctx.Err() } out := make(chan types.StreamResponse, len(m.chunks)) for _, chunk := range m.chunks { out <- chunk } close(out) return out, nil } func collectDigitalHumanTrial(t *testing.T, stream <-chan DigitalHumanTrialEvent) []DigitalHumanTrialEvent { t.Helper() var events []DigitalHumanTrialEvent for { select { case event, ok := <-stream: if !ok { return events } events = append(events, event) case <-time.After(5 * time.Second): t.Fatal("trial stream did not finish") } } } func TestDigitalHumanTrialUsesSelectedDraftWithoutPublishing(t *testing.T) { svc, ctx, human, _, _ := setupDigitalHumanWeChat(t) svc.slots = make(chan struct{}, 2) model := &digitalHumanTrialChat{chunks: []types.StreamResponse{ {ResponseType: types.ResponseTypeThinking, Content: "private reasoning"}, {ResponseType: types.ResponseTypeAnswer, Content: "第一句。"}, {ResponseType: types.ResponseTypeAnswer, Content: "第二句。", Done: true, FinishReason: "stop"}, }} svc.models = &digitalHumanTrialModelService{model: model, status: types.ModelStatusActive} scope, _ := digitalHumanScope(ctx) require.NoError(t, svc.repo.Review(ctx, scope, human.ID, 1, true, "approved")) active, err := svc.repo.Get(ctx, scope, human.ID) require.NoError(t, err) require.NoError(t, svc.repo.Update(ctx, scope, human.ID, active.Revision, active.Name, "", "", "configured-model", map[string]string{"tone": "selected draft style"})) before, err := svc.repo.Get(ctx, scope, human.ID) require.NoError(t, err) stream, err := svc.Trial(ctx, human.ID, DigitalHumanTrialInput{Version: 2, Theme: "教学主题", Advice: "先举例,再总结"}) require.NoError(t, err) events := collectDigitalHumanTrial(t, stream) require.Equal(t, "progress", events[0].Type) require.Equal(t, "complete", events[len(events)-1].Type) require.Equal(t, 100, events[len(events)-1].Progress) var output strings.Builder for _, event := range events { require.Equal(t, 2, event.Version) require.Equal(t, "configured-model", event.ModelID) output.WriteString(event.Content) } require.Equal(t, "第一句。第二句。", output.String()) var payload map[string]any require.NoError(t, json.Unmarshal([]byte(model.messages[1].Content), &payload)) require.Equal(t, float64(2), payload["selected_version"]) require.Equal(t, map[string]any{"tone": "selected draft style"}, payload["expression_profile"]) require.Equal(t, "先举例,再总结", payload["writing_advice"]) require.Equal(t, "digital_human_trial", model.purpose) require.Equal(t, "owner", model.owner) require.Equal(t, uint64(7), model.tenant) after, err := svc.repo.Get(ctx, scope, human.ID) require.NoError(t, err) require.Equal(t, before, after) require.Equal(t, 1, after.ActiveVersion) require.Equal(t, 1, after.Bindings[0].AppliedVersion) require.Empty(t, svc.slots) } func TestDigitalHumanTrialRejectsOtherOwnerTenantAndInvalidInput(t *testing.T) { svc, ctx, human, _, _ := setupDigitalHumanWeChat(t) svc.slots = make(chan struct{}, 2) svc.models = &digitalHumanTrialModelService{model: &digitalHumanTrialChat{}, status: types.ModelStatusActive} input := DigitalHumanTrialInput{Version: 1, Theme: "topic"} for _, other := range []context.Context{ context.WithValue(ctx, types.UserIDContextKey, "other"), context.WithValue(ctx, types.TenantIDContextKey, uint64(8)), context.WithValue(context.WithValue(ctx, types.UserIDContextKey, "admin"), types.TenantRoleContextKey, types.TenantRoleAdmin), accesspolicy.WithSnapshot(context.WithValue(ctx, types.UserIDContextKey, "shared-viewer"), &accesspolicy.Snapshot{Grants: map[string]bool{"digital-humans:" + human.ID: true}}), } { _, err := svc.Trial(other, human.ID, input) require.ErrorIs(t, err, types.ErrDigitalHumanNotFound) } _, err := svc.Trial(accesspolicy.WithSnapshot(ctx, &accesspolicy.Snapshot{Restricted: true, Permissions: map[string]bool{"digital-humans.read": true}}), human.ID, input) require.ErrorContains(t, err, "没有权限") for _, invalid := range []DigitalHumanTrialInput{ {Version: 0, Theme: "topic"}, {Version: 1}, {Version: 1, Theme: strings.Repeat("文", 1001)}, {Version: 1, Theme: "topic", Advice: strings.Repeat("文", 6001)}, {Version: 999, Theme: "topic"}, } { _, err := svc.Trial(ctx, human.ID, invalid) require.Error(t, err) } require.Empty(t, svc.slots) } func TestDigitalHumanTrialReportsIncompleteAndModelFailures(t *testing.T) { for _, test := range []struct { name string chunks []types.StreamResponse err error }{ {"provider_failure", nil, errors.New("secret-provider-message")}, {"empty", []types.StreamResponse{{Done: true}}, nil}, {"missing_done", []types.StreamResponse{{ResponseType: types.ResponseTypeAnswer, Content: "partial"}}, nil}, {"truncated", []types.StreamResponse{{ResponseType: types.ResponseTypeAnswer, Content: "partial", Done: true, FinishReason: "length"}}, nil}, {"settlement_failure", []types.StreamResponse{{Done: true, FinishReason: types.FinishReasonIncomplete}}, nil}, } { t.Run(test.name, func(t *testing.T) { svc := &DigitalHumanService{slots: make(chan struct{}, 1)} svc.slots <- struct{}{} out := make(chan DigitalHumanTrialEvent) go svc.runDigitalHumanTrial(context.Background(), out, &digitalHumanTrialChat{chunks: test.chunks, err: test.err}, 1, "configured", "{}") events := collectDigitalHumanTrial(t, out) require.Equal(t, "error", events[len(events)-1].Type) require.NotContains(t, events[len(events)-1].Error, "secret-provider-message") require.Empty(t, svc.slots) }) } } func TestDigitalHumanTrialDisconnectCancelsProviderAndReleasesSlot(t *testing.T) { svc := &DigitalHumanService{slots: make(chan struct{}, 1)} svc.slots <- struct{}{} model := &digitalHumanTrialChat{block: true, cancelled: make(chan struct{})} ctx, cancel := context.WithCancel(context.Background()) out := make(chan DigitalHumanTrialEvent) go svc.runDigitalHumanTrial(ctx, out, model, 1, "configured", "{}") require.Equal(t, "progress", (<-out).Type) cancel() collectDigitalHumanTrial(t, out) select { case <-model.cancelled: case <-time.After(time.Second): t.Fatal("provider was not cancelled") } require.Empty(t, svc.slots) }