xinmi/internal/application/service/resumable_upload_test.go

package service

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"mime/multipart"
	"os"
	"path/filepath"
	"strings"
	"testing"
	"unicode/utf8"

	"github.com/Tencent/WeKnora/internal/application/repository"
	"github.com/Tencent/WeKnora/internal/resumable"
	"github.com/Tencent/WeKnora/internal/types"
	"github.com/Tencent/WeKnora/internal/types/interfaces"
	"github.com/hibiken/asynq"
	"github.com/stretchr/testify/require"
)

func TestResumableTextPartsPreserveUTF8AndBytes(t *testing.T) {
	input := []byte(strings.Repeat("原文🙂\n", uploadTextPartBytes/4))
	path := filepath.Join(t.TempDir(), "source.txt")
	require.NoError(t, os.WriteFile(path, input, 0600))
	output := []byte{}
	offset := int64(0)
	for offset < int64(len(input)) {
		part, next, name, e := readUploadPart(path, "source.txt", offset, int64(len(input)))
		require.NoError(t, e)
		require.True(t, utf8.Valid(part))
		require.LessOrEqual(t, len(part), uploadTextPartBytes)
		require.Greater(t, next, offset)
		require.Contains(t, name, ".part-")
		output = append(output, part...)
		offset = next
	}
	require.Equal(t, input, output)
	_, _, _, e := readUploadPart("missing", "large.pdf", 0, uploadBinaryParseBytes+1)
	require.ErrorContains(t, e, "32 MiB")
}

type resumableQueueRecoveryFake struct {
	interfaces.KnowledgeService
	doc              *types.Knowledge
	missing          bool
	repairs, creates int
}

func (s *resumableQueueRecoveryFake) GetKnowledgeByID(context.Context, string) (*types.Knowledge, error) {
	if s.missing {
		return nil, repository.ErrKnowledgeNotFound
	}
	return s.doc, nil
}
func (s *resumableQueueRecoveryFake) CreateKnowledgeFromFile(ctx context.Context, _ string, _ *multipart.FileHeader, _ map[string]string, _ *bool, _ string, _ []string, _ string, _ *types.KnowledgeProcessOverrides) (*types.Knowledge, error) {
	s.creates++
	s.missing = false
	s.doc.ID = ctx.Value(resumableKnowledgeIDKey{}).(string)
	return s.doc, types.NewDuplicateFileError(s.doc)
}
func (s *resumableQueueRecoveryFake) EnsureResumableKnowledgeQueued(_ context.Context, doc *types.Knowledge, _ []string) (*types.Knowledge, error) {
	s.repairs++
	doc.ParseStatus = "completed"
	return doc, nil
}

type resumableRepairTask struct {
	task   *asynq.Task
	opts   []asynq.Option
	err    error
	active bool
}

func (s *resumableRepairTask) HasActiveKnowledgeTask(string) bool { return s.active }

func (s *resumableRepairTask) Enqueue(task *asynq.Task, opts ...asynq.Option) (*asynq.TaskInfo, error) {
	s.task = task
	s.opts = opts
	return &asynq.TaskInfo{ID: "repair"}, s.err
}

func TestResumableNativeRepairUsesStableQueueID(t *testing.T) {
	for _, conflict := range []error{nil, asynq.ErrTaskIDConflict} {
		task := &resumableRepairTask{err: conflict}
		s := &knowledgeService{repo: &createKnowledgeFileRepoStub{}, task: task, kbService: &createKnowledgeFileKBServiceStub{kb: &types.KnowledgeBase{ID: "kb"}}}
		doc := &types.Knowledge{ID: "reserved-doc", TenantID: 1, KnowledgeBaseID: "kb", ParseStatus: "pending", FilePath: "stored/source", FileName: "source.txt", FileType: "txt"}
		result, e := s.EnsureResumableKnowledgeQueued(newCreateKnowledgeFileContext(), doc, nil)
		require.NoError(t, e)
		require.Same(t, doc, result)
		require.NotNil(t, task.task)
		stableID := false
		for _, option := range task.opts {
			if option.Type() == asynq.TaskIDOpt {
				stableID = option.Value() == "resumable:reserved-doc"
			}
		}
		require.True(t, stableID)
		require.Contains(t, string(task.task.Payload()), `"knowledge_id":"reserved-doc"`)
	}
}

func TestResumableNativeRepairDoesNotReparseActiveFailure(t *testing.T) {
	s := &knowledgeService{taskInspector: fakeTaskInspector{queued: map[string]bool{"reserved-doc": true}}}
	doc := &types.Knowledge{ID: "reserved-doc", ParseStatus: "failed"}
	result, e := s.EnsureResumableKnowledgeQueued(context.Background(), doc, nil)
	require.NoError(t, e)
	require.Same(t, doc, result)
	s = &knowledgeService{task: &resumableRepairTask{active: true}}
	result, e = s.EnsureResumableKnowledgeQueued(context.Background(), doc, nil)
	require.NoError(t, e)
	require.Same(t, doc, result)
	ctx, cancel := context.WithCancel(context.Background())
	cancel()
	_, e = s.EnsureResumableKnowledgeQueued(ctx, doc, nil)
	require.ErrorIs(t, e, context.Canceled)
}

func TestResumableNativeCrashBeforeEnqueueRepairsExactDocument(t *testing.T) {
	for _, concurrentInsert := range []bool{false, true} {
		t.Run(fmt.Sprint(concurrentInsert), func(t *testing.T) {
			path := filepath.Join(t.TempDir(), "source.txt")
			require.NoError(t, os.WriteFile(path, []byte("source"), 0600))
			row := &resumable.Upload{ID: "upload", TenantID: 1, TargetID: "kb", Name: "source.txt", Size: 6, Options: types.JSON(`{}`), Result: types.JSON(`{}`)}
			id := resumableDocumentID(row)
			fake := &resumableQueueRecoveryFake{missing: concurrentInsert, doc: &types.Knowledge{ID: id, TenantID: 1, KnowledgeBaseID: "kb", ParseStatus: "pending"}}
			offset, result, done, e := ingestResumableKnowledge(context.Background(), fake, row, path)
			require.NoError(t, e)
			require.True(t, done)
			require.EqualValues(t, 6, offset)
			require.Contains(t, string(result), id)
			require.Equal(t, 1, fake.repairs)
			if concurrentInsert {
				require.Equal(t, 1, fake.creates)
			} else {
				require.Zero(t, fake.creates)
			}
		})
	}
}

type resumableCollisionRepo struct{ createKnowledgeFileRepoStub }

func (r *resumableCollisionRepo) GetKnowledgeByID(context.Context, uint64, string) (*types.Knowledge, error) {
	copy := *r.createdKnowledge
	copy.FilePath = "existing-object"
	return &copy, nil
}
func TestResumableNativePrimaryKeyCollisionPreservesWinningFile(t *testing.T) {
	repo := &resumableCollisionRepo{createKnowledgeFileRepoStub: createKnowledgeFileRepoStub{createErr: errors.New("unique key conflict")}}
	storage := &createKnowledgeFileServiceStub{}
	s := &knowledgeService{repo: repo, kbService: &createKnowledgeFileKBServiceStub{kb: &types.KnowledgeBase{ID: "kb-1"}}, fileSvc: storage}
	ctx := context.WithValue(newCreateKnowledgeFileContext(), resumableKnowledgeIDKey{}, "a1c342cd-1c1f-59b9-a9fa-2e883dfb4a3e")
	doc, e := s.CreateKnowledgeFromFile(ctx, "kb-1", newMultipartFileHeader(t, "source.txt", "source"), nil, nil, "", nil, "", nil)
	var duplicate *types.DuplicateKnowledgeError
	require.ErrorAs(t, e, &duplicate)
	require.Equal(t, "a1c342cd-1c1f-59b9-a9fa-2e883dfb4a3e", doc.ID)
	require.Equal(t, "existing-object", doc.FilePath)
	require.Equal(t, 1, storage.deleteCalls)
	require.NotEqual(t, doc.FilePath, storage.deletedPath)
}

func TestResumableMaterialReadsMiddleAndEndWithoutSampling(t *testing.T) {
	data := []byte("BEGIN " + strings.Repeat("本人原文\n", 15000) + " UNIQUE MIDDLE " + strings.Repeat("后半原文\n", 15000) + " END")
	source, selection, e := readFullMaterial("large.txt", bytes.NewReader(data), int64(len(data)))
	require.NoError(t, e)
	require.Equal(t, "full", selection)
	require.Equal(t, string(data), source.Text)
	source, selection, e = readFullMaterial("small.txt", strings.NewReader("完整原文"), int64(len("完整原文")))
	require.NoError(t, e)
	require.Equal(t, "full", selection)
	require.Equal(t, "完整原文", source.Text)
	_, _, e = readFullMaterial("bad.txt", bytes.NewReader([]byte{0xff}), 1)
	require.Error(t, e)
	_, _, e = readFullMaterial("too-large.txt", strings.NewReader(""), int64(digitalHumanSourceLimit*utf8.UTFMax+4))
	require.ErrorContains(t, e, "24万字")
	_, _, e = readFullMaterial("too-long.txt", strings.NewReader(strings.Repeat("x", digitalHumanSourceLimit+1)), digitalHumanSourceLimit+1)
	require.ErrorContains(t, e, "24万字")
	_, _, e = readFullMaterial("truncated.txt", strings.NewReader("partial"), 100)
	require.Error(t, e)
}