xinmi/internal/application/service/digital_human_trial_test.go

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)
}