WeKnora deployment source with the reviewed wiki identity, summary and recovery fixes

internal/application/service/digital_human_trial_test.go

@@ -0,0 +1,198 @@
+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)
+}