xinmi/internal/application/service/wiki_call_deadline_test.go

package service

import (
	"context"
	"errors"
	"testing"
	"time"

	"github.com/Tencent/WeKnora/internal/models/chat"
	"github.com/Tencent/WeKnora/internal/types"
	"github.com/stretchr/testify/require"
)

func TestWikiCallDeadlineBoundsEachAttemptWithinCheckpointSlice(t *testing.T) {
	ctx, cancel := context.WithTimeout(context.Background(), wikiCheckpointSliceTimeout)
	defer cancel()
	var calls []context.Context
	model := &wikiScriptedChat{id: t.Name(), call: func(callCtx context.Context, _ []chat.Message) (*types.ChatResponse, error) {
		calls = append(calls, callCtx)
		deadline, ok := callCtx.Deadline()
		require.True(t, ok)
		require.InDelta(t, wikiLLMCallTimeout.Seconds(), time.Until(deadline).Seconds(), 1)
		if len(calls) == 1 {
			return nil, context.DeadlineExceeded
		}
		require.ErrorIs(t, calls[0].Err(), context.Canceled, "finished attempt must release its timer")
		require.NoError(t, callCtx.Err(), "retry must receive a fresh deadline")
		return &types.ChatResponse{Content: "complete", FinishReason: "stop"}, nil
	}}

	output, err := (&wikiIngestService{}).generateWithTemplate(ctx, model, "{{.Content}}", map[string]string{"Content": "source"})
	require.NoError(t, err)
	require.Equal(t, "complete", output)
	require.Len(t, calls, 2)
	require.NoError(t, ctx.Err(), "single-call failure must not cancel the checkpoint slice")
	require.ErrorIs(t, calls[1].Err(), context.Canceled)
}

func TestWikiCallDeadlineHonorsRecoveryDeadlineWithoutRetry(t *testing.T) {
	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Millisecond)
	defer cancel()
	parentDeadline, _ := ctx.Deadline()
	calls := 0
	model := &wikiScriptedChat{id: t.Name(), call: func(callCtx context.Context, _ []chat.Message) (*types.ChatResponse, error) {
		calls++
		deadline, ok := callCtx.Deadline()
		require.True(t, ok)
		require.Equal(t, parentDeadline, deadline)
		<-callCtx.Done()
		return nil, callCtx.Err()
	}}

	output, err := (&wikiIngestService{}).generateWithTemplate(ctx, model, "{{.Content}}", map[string]string{"Content": "source"})
	require.Empty(t, output)
	require.True(t, errors.Is(err, context.DeadlineExceeded))
	require.Equal(t, 1, calls, "expired recovery/task budget must not start another model call")
}