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