xinmi/internal/application/service/knowledge_batch_reparse_test.go

package service

import (
	"context"
	"encoding/json"
	"errors"
	"testing"
	"time"

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

type reparseFailureKnowledgeRepo struct {
	interfaces.KnowledgeRepository
	knowledge   *types.Knowledge
	updateCalls int
}

func (r *reparseFailureKnowledgeRepo) GetKnowledgeByID(
	_ context.Context,
	_ uint64,
	_ string,
) (*types.Knowledge, error) {
	return r.knowledge, nil
}

func (r *reparseFailureKnowledgeRepo) UpdateKnowledge(
	_ context.Context,
	_ *types.Knowledge,
) error {
	r.updateCalls++
	return nil
}

func (r *reparseFailureKnowledgeRepo) UpdateKnowledgeColumn(
	_ context.Context,
	_ string,
	_ string,
	_ interface{},
) error {
	return nil
}

type reparseFailureKBService struct {
	interfaces.KnowledgeBaseService
	kb *types.KnowledgeBase
}

func (s *reparseFailureKBService) GetKnowledgeBaseByID(
	_ context.Context,
	_ string,
) (*types.KnowledgeBase, error) {
	return s.kb, nil
}

type failingReparseTaskEnqueuer struct {
	err error
}

func (e failingReparseTaskEnqueuer) Enqueue(
	_ *asynq.Task,
	_ ...asynq.Option,
) (*asynq.TaskInfo, error) {
	return nil, e.err
}

func TestReparseKnowledgeManualEnqueueFailureIsVisible(t *testing.T) {
	enqueueErr := errors.New("queue unavailable")
	knowledge := &types.Knowledge{
		ID:              "knowledge-1",
		TenantID:        7,
		KnowledgeBaseID: "kb-1",
		Type:            types.KnowledgeTypeManual,
		ParseStatus:     types.ParseStatusCompleted,
		EnableStatus:    "enabled",
	}
	require.NoError(t, knowledge.SetManualMetadata(
		types.NewManualKnowledgeMetadata("# content", types.ManualKnowledgeStatusPublish, 1),
	))
	repo := &reparseFailureKnowledgeRepo{knowledge: knowledge}
	svc := &knowledgeService{
		repo:      repo,
		kbService: &reparseFailureKBService{kb: &types.KnowledgeBase{ID: "kb-1"}},
		task:      failingReparseTaskEnqueuer{err: enqueueErr},
	}
	ctx := context.WithValue(context.Background(), types.TenantIDContextKey, uint64(7))

	got, err := svc.ReparseKnowledge(ctx, knowledge.ID, nil)

	require.Error(t, err)
	require.Same(t, knowledge, got)
	require.Equal(t, types.ParseStatusFailed, knowledge.ParseStatus)
	require.Equal(t, "disabled", knowledge.EnableStatus)
	require.Equal(t, "Failed to enqueue processing task", knowledge.ErrorMessage)
	require.GreaterOrEqual(t, repo.updateCalls, 2, "pending and failed states must both be persisted")
}

type reparseOptionsTaskEnqueuer struct {
	task *asynq.Task
	opts []asynq.Option
}

func (e *reparseOptionsTaskEnqueuer) Enqueue(task *asynq.Task, opts ...asynq.Option) (*asynq.TaskInfo, error) {
	e.task = task
	e.opts = opts
	return &asynq.TaskInfo{ID: "reparse-task", Queue: types.QueueDefault}, nil
}

type reparseOptionsChunkService struct {
	parentChildChunkService
}

func (reparseOptionsChunkService) GetRepository() interfaces.ChunkRepository {
	return kbCleanupChunkRepo{}
}

func TestReparseKnowledgePassesDocumentOptionsToEnqueue(t *testing.T) {
	for _, tc := range []struct {
		name     string
		kind     string
		filePath string
		source   string
	}{
		{name: "file", kind: "file", filePath: "stored/document.txt"},
		{name: "file URL", kind: "file_url", source: "https://example.com/document.txt"},
		{name: "web URL", kind: "url", source: "https://example.com/document"},
	} {
		t.Run(tc.name, func(t *testing.T) {
			knowledge := &types.Knowledge{
				ID: "knowledge-1", TenantID: 7, KnowledgeBaseID: "kb-1",
				Type: tc.kind, FilePath: tc.filePath, Source: tc.source,
				FileName: "document.txt", FileType: "txt", ParseStatus: types.ParseStatusCompleted,
			}
			queue := &reparseOptionsTaskEnqueuer{}
			svc := &knowledgeService{
				config: &config.Config{KnowledgeBase: &config.KnowledgeBaseConfig{
					DocumentProcessTimeout: 90 * time.Minute,
				}},
				repo:         &reparseFailureKnowledgeRepo{knowledge: knowledge},
				kbService:    &reparseFailureKBService{kb: &types.KnowledgeBase{ID: "kb-1"}},
				chunkService: &reparseOptionsChunkService{},
				graphEngine:  parentChildGraphRepo{},
				task:         queue,
			}
			ctx := context.WithValue(context.Background(), types.TenantIDContextKey, uint64(7))
			ctx = context.WithValue(ctx, types.TenantInfoContextKey, &types.Tenant{ID: 7})

			_, err := svc.ReparseKnowledge(ctx, knowledge.ID, nil)

			require.NoError(t, err)
			require.NotNil(t, queue.task)
			require.Equal(t, types.TypeDocumentProcess, queue.task.Type())
			var payload types.DocumentProcessPayload
			require.NoError(t, json.Unmarshal(queue.task.Payload(), &payload))
			require.Equal(t, knowledge.ID, payload.KnowledgeID)
			queueName, timeout, maxRetry := parseDocumentProcessOpts(t, queue.opts)
			require.Equal(t, types.QueueDefault, queueName)
			require.Equal(t, 90*time.Minute, timeout)
			require.NotNil(t, maxRetry, "Lite only receives options passed directly to Enqueue")
			require.Equal(t, 3, *maxRetry)
		})
	}
}

func TestRunKnowledgeListReparseSubmissionsReportsPartialFailure(t *testing.T) {
	firstErr := errors.New("first failed")
	secondErr := errors.New("second failed")
	var attempted []string

	outcome, err := runKnowledgeListReparseSubmissions(
		[]string{"ok-1", "bad-1", "ok-2", "bad-2"},
		func(id string) error {
			attempted = append(attempted, id)
			switch id {
			case "bad-1":
				return firstErr
			case "bad-2":
				return secondErr
			default:
				return nil
			}
		},
	)

	require.Equal(t, []string{"ok-1", "bad-1", "ok-2", "bad-2"}, attempted)
	require.Equal(t, knowledgeListReparseOutcome{Submitted: 2, Failed: 2}, outcome)
	require.ErrorIs(t, err, asynq.SkipRetry)
	require.ErrorIs(t, err, firstErr)
	require.ErrorIs(t, err, secondErr)
	require.ErrorContains(t, err, "knowledge bad-1")
	require.ErrorContains(t, err, "knowledge bad-2")
}

func TestRunKnowledgeListReparseSubmissionsSucceeds(t *testing.T) {
	outcome, err := runKnowledgeListReparseSubmissions(
		[]string{"knowledge-1", "knowledge-2"},
		func(string) error { return nil },
	)

	require.NoError(t, err)
	require.Equal(t, knowledgeListReparseOutcome{Submitted: 2}, outcome)
}