mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-07-26 18:33:30 +08:00
047c714ffb
### Summary 1. Remove dead code (replaced by builtin ingestion pipeline) 2. Refactor (move document parsing progress from http api into ingestion executor)
877 lines
28 KiB
Go
877 lines
28 KiB
Go
package service
|
|
|
|
import (
|
|
"errors"
|
|
"testing"
|
|
|
|
"ragflow/internal/common"
|
|
"ragflow/internal/dao"
|
|
"ragflow/internal/entity"
|
|
)
|
|
|
|
type recordingTaskPublisher struct {
|
|
subject string
|
|
messages []common.TaskMessage
|
|
err error
|
|
}
|
|
|
|
func (p *recordingTaskPublisher) PublishTaskMessage(subject string, msg common.TaskMessage) error {
|
|
p.subject = subject
|
|
p.messages = append(p.messages, msg)
|
|
return p.err
|
|
}
|
|
|
|
func TestIngestionTaskServiceCreateForDocumentsPublishesTaskMessages(t *testing.T) {
|
|
db := setupServiceTestDB(t)
|
|
pushServiceDB(t, db)
|
|
insertTestKB(t, "kb-1", "tenant-1", 1, 0, 0)
|
|
insertTestDoc(t, "doc-1", "kb-1", 0, 0)
|
|
|
|
publisher := &recordingTaskPublisher{}
|
|
svc := NewIngestionTaskService()
|
|
svc.taskPublisher = publisher
|
|
|
|
resp, err := svc.CreateForDocuments("kb-1", "user-1", []string{"doc-1"})
|
|
if err != nil {
|
|
t.Fatalf("CreateForDocuments failed: %v", err)
|
|
}
|
|
if len(resp) != 1 {
|
|
t.Fatalf("expected 1 response, got %d", len(resp))
|
|
}
|
|
if publisher.subject != "tasks.RAGFLOW" {
|
|
t.Fatalf("subject = %q, want %q", publisher.subject, "tasks.RAGFLOW")
|
|
}
|
|
if len(publisher.messages) != 1 {
|
|
t.Fatalf("expected 1 published message, got %d", len(publisher.messages))
|
|
}
|
|
msg := publisher.messages[0]
|
|
if msg.TaskType != common.TaskTypeIngestionTask {
|
|
t.Fatalf("task type = %q, want %q", msg.TaskType, common.TaskTypeIngestionTask)
|
|
}
|
|
task, err := dao.NewIngestionTaskDAO().GetByID(msg.TaskID)
|
|
if err != nil {
|
|
t.Fatalf("load task: %v", err)
|
|
}
|
|
if task.DocumentID != "doc-1" || task.DatasetID != "kb-1" || task.UserID != "user-1" {
|
|
t.Fatalf("unexpected task: %+v", task)
|
|
}
|
|
}
|
|
|
|
func TestIngestionTaskServiceListByUserFiltersDataset(t *testing.T) {
|
|
db := setupServiceTestDB(t)
|
|
pushServiceDB(t, db)
|
|
insertTestIngestionTask(t, "task-1", "user-1", "doc-1", "kb-1")
|
|
insertTestIngestionTask(t, "task-2", "user-1", "doc-2", "kb-2")
|
|
insertTestIngestionTask(t, "task-3", "user-2", "doc-3", "kb-1")
|
|
|
|
svc := NewIngestionTaskService()
|
|
datasetID := "kb-1"
|
|
tasks, err := svc.ListByUser("user-1", &datasetID, 0, 0)
|
|
if err != nil {
|
|
t.Fatalf("ListByUser failed: %v", err)
|
|
}
|
|
if len(tasks) != 1 {
|
|
t.Fatalf("expected 1 task, got %d", len(tasks))
|
|
}
|
|
if tasks[0].ID != "task-1" {
|
|
t.Fatalf("task ID = %q, want %q", tasks[0].ID, "task-1")
|
|
}
|
|
}
|
|
|
|
func TestIngestionTaskServiceRequestStopManyStopsOwnedTasks(t *testing.T) {
|
|
db := setupServiceTestDB(t)
|
|
pushServiceDB(t, db)
|
|
insertTestIngestionTask(t, "task-1", "user-1", "doc-1", "kb-1")
|
|
|
|
userID := "user-1"
|
|
svc := NewIngestionTaskService()
|
|
tasks, err := svc.RequestStopMany([]string{"task-1"}, &userID)
|
|
if err != nil {
|
|
t.Fatalf("RequestStopMany failed: %v", err)
|
|
}
|
|
if len(tasks) != 1 {
|
|
t.Fatalf("expected 1 task response, got %d", len(tasks))
|
|
}
|
|
if tasks[0].Status != common.STOPPED {
|
|
t.Fatalf("status = %q, want %q", tasks[0].Status, common.STOPPED)
|
|
}
|
|
}
|
|
|
|
func TestIngestionTaskServiceRequestStopManyRejectsOtherUsersTask(t *testing.T) {
|
|
db := setupServiceTestDB(t)
|
|
pushServiceDB(t, db)
|
|
insertTestIngestionTask(t, "task-1", "user-1", "doc-1", "kb-1")
|
|
|
|
userID := "user-2"
|
|
svc := NewIngestionTaskService()
|
|
if _, err := svc.RequestStopMany([]string{"task-1"}, &userID); err == nil {
|
|
t.Fatal("expected RequestStopMany to reject non-owner")
|
|
}
|
|
task, err := dao.NewIngestionTaskDAO().GetByID("task-1")
|
|
if err != nil {
|
|
t.Fatalf("load task: %v", err)
|
|
}
|
|
if task.Status != common.CREATED {
|
|
t.Fatalf("status = %q, want %q", task.Status, common.CREATED)
|
|
}
|
|
}
|
|
|
|
func TestIngestionTaskServiceRequestStopManyAllowsAdmin(t *testing.T) {
|
|
db := setupServiceTestDB(t)
|
|
pushServiceDB(t, db)
|
|
insertTestIngestionTask(t, "task-1", "user-1", "doc-1", "kb-1")
|
|
|
|
svc := NewIngestionTaskService()
|
|
tasks, err := svc.RequestStopMany([]string{"task-1"}, nil)
|
|
if err != nil {
|
|
t.Fatalf("RequestStopMany admin failed: %v", err)
|
|
}
|
|
if len(tasks) != 1 {
|
|
t.Fatalf("expected 1 task response, got %d", len(tasks))
|
|
}
|
|
if tasks[0].Status != common.STOPPED {
|
|
t.Fatalf("status = %q, want %q", tasks[0].Status, common.STOPPED)
|
|
}
|
|
}
|
|
|
|
func TestIngestionTaskServiceRemoveManyRemovesOwnedTasks(t *testing.T) {
|
|
db := setupServiceTestDB(t)
|
|
pushServiceDB(t, db)
|
|
insertTestIngestionTask(t, "task-1", "user-1", "doc-1", "kb-1")
|
|
|
|
userID := "user-1"
|
|
svc := NewIngestionTaskService()
|
|
result, err := svc.RemoveMany([]string{"task-1"}, &userID)
|
|
if err != nil {
|
|
t.Fatalf("RemoveMany failed: %v", err)
|
|
}
|
|
if len(result) != 1 || result[0]["remove"] != "success" {
|
|
t.Fatalf("unexpected remove result: %+v", result)
|
|
}
|
|
if _, err := dao.NewIngestionTaskDAO().GetByID("task-1"); err == nil {
|
|
t.Fatal("task should be removed")
|
|
}
|
|
}
|
|
|
|
func TestIngestionTaskServiceListAllForAdminIncludesRunAndUserEmail(t *testing.T) {
|
|
db := setupServiceTestDB(t)
|
|
pushServiceDB(t, db)
|
|
status := "1"
|
|
if err := dao.DB.Create(&entity.User{
|
|
ID: "user-1",
|
|
Email: "user-1@test.com",
|
|
Nickname: "user-1",
|
|
IsAuthenticated: "1",
|
|
IsActive: "1",
|
|
IsAnonymous: "0",
|
|
Status: &status,
|
|
}).Error; err != nil {
|
|
t.Fatalf("insert user: %v", err)
|
|
}
|
|
insertTestIngestionTask(t, "task-1", "user-1", "doc-1", "kb-1")
|
|
if err := dao.DB.Create(&entity.IngestionTaskLog{
|
|
TaskID: "task-1",
|
|
Checkpoint: entity.JSONMap{"run_count": 3},
|
|
}).Error; err != nil {
|
|
t.Fatalf("insert task log: %v", err)
|
|
}
|
|
|
|
svc := NewIngestionTaskService()
|
|
tasks, err := svc.ListAllForAdmin()
|
|
if err != nil {
|
|
t.Fatalf("ListAllForAdmin failed: %v", err)
|
|
}
|
|
if len(tasks) != 1 {
|
|
t.Fatalf("expected 1 task, got %d", len(tasks))
|
|
}
|
|
if tasks[0]["user"] != "user-1@test.com" {
|
|
t.Fatalf("user = %v, want user-1@test.com", tasks[0]["user"])
|
|
}
|
|
if tasks[0]["run_count"] != 3 {
|
|
t.Fatalf("run_count = %v, want 3", tasks[0]["run_count"])
|
|
}
|
|
if tasks[0]["component_total"] != 0 {
|
|
t.Fatalf("component_total = %v, want 0", tasks[0]["component_total"])
|
|
}
|
|
if tasks[0]["component_done"] != 0 {
|
|
t.Fatalf("component_done = %v, want 0", tasks[0]["component_done"])
|
|
}
|
|
}
|
|
|
|
func TestIngestionTaskServiceStartRunningTransitionsCreatedTask(t *testing.T) {
|
|
db := setupServiceTestDB(t)
|
|
pushServiceDB(t, db)
|
|
insertTestIngestionTask(t, "task-1", "user-1", "doc-1", "kb-1")
|
|
|
|
svc := NewIngestionTaskService()
|
|
task, err := svc.StartRunning("task-1")
|
|
if err != nil {
|
|
t.Fatalf("StartRunning failed: %v", err)
|
|
}
|
|
if task.Status != common.RUNNING {
|
|
t.Fatalf("status = %q, want %q", task.Status, common.RUNNING)
|
|
}
|
|
}
|
|
|
|
// TestStartRunningMarksDocumentRunning locks in that starting a CREATED task
|
|
// mirrors the transition to its document: run=RUNNING and progress counters
|
|
// reset, with a fresh process_begin_at. The document bookkeeping is owned by
|
|
// the task-lifecycle transition, not the ingestion worker's execution path.
|
|
func TestStartRunningMarksDocumentRunning(t *testing.T) {
|
|
db := setupServiceTestDB(t)
|
|
pushServiceDB(t, db)
|
|
insertTestKB(t, "kb-1", "tenant-1", 1, 0, 0)
|
|
insertTestDoc(t, "doc-1", "kb-1", 100, 10)
|
|
insertTestIngestionTask(t, "task-1", "user-1", "doc-1", "kb-1")
|
|
|
|
// Seed the document as a partially-processed, non-RUNNING state that the
|
|
// start transition must clobber.
|
|
if err := db.Model(&entity.Document{}).Where("id = ?", "doc-1").
|
|
Updates(map[string]interface{}{
|
|
"run": string(entity.TaskStatusDone),
|
|
"progress": float64(0.5),
|
|
"progress_msg": "partial",
|
|
}).Error; err != nil {
|
|
t.Fatalf("seed document: %v", err)
|
|
}
|
|
|
|
svc := NewIngestionTaskService()
|
|
if _, err := svc.StartRunning("task-1"); err != nil {
|
|
t.Fatalf("StartRunning failed: %v", err)
|
|
}
|
|
|
|
var doc entity.Document
|
|
if err := db.Where("id = ?", "doc-1").First(&doc).Error; err != nil {
|
|
t.Fatalf("reload document: %v", err)
|
|
}
|
|
if doc.Run == nil || *doc.Run != string(entity.TaskStatusRunning) {
|
|
t.Fatalf("run = %v, want RUNNING(%q)", doc.Run, string(entity.TaskStatusRunning))
|
|
}
|
|
if doc.Progress != 0 {
|
|
t.Fatalf("progress = %f, want 0", doc.Progress)
|
|
}
|
|
if doc.ChunkNum != 0 {
|
|
t.Fatalf("chunk_num = %d, want 0", doc.ChunkNum)
|
|
}
|
|
if doc.TokenNum != 0 {
|
|
t.Fatalf("token_num = %d, want 0", doc.TokenNum)
|
|
}
|
|
if doc.ProgressMsg != nil && *doc.ProgressMsg != "" {
|
|
t.Fatalf("progress_msg = %q, want empty", *doc.ProgressMsg)
|
|
}
|
|
if doc.ProcessBeginAt == nil || doc.ProcessBeginAt.IsZero() {
|
|
t.Fatal("process_begin_at not set")
|
|
}
|
|
}
|
|
|
|
// TestStartRunningLeavesTerminalDocumentUntouched locks in the no-resurrection
|
|
// invariant: a task already in a terminal status is returned as-is by
|
|
// StartRunning, and its document's finished run status/counters are not reset.
|
|
func TestStartRunningLeavesTerminalDocumentUntouched(t *testing.T) {
|
|
db := setupServiceTestDB(t)
|
|
pushServiceDB(t, db)
|
|
insertTestKB(t, "kb-1", "tenant-1", 1, 0, 0)
|
|
insertTestDoc(t, "doc-1", "kb-1", 100, 10)
|
|
insertTestIngestionTask(t, "task-1", "user-1", "doc-1", "kb-1")
|
|
|
|
finishedRun := string(entity.TaskStatusDone)
|
|
if err := db.Model(&entity.Document{}).Where("id = ?", "doc-1").
|
|
Updates(map[string]interface{}{
|
|
"run": finishedRun,
|
|
"progress": float64(1.0),
|
|
"progress_msg": "done",
|
|
}).Error; err != nil {
|
|
t.Fatalf("seed document: %v", err)
|
|
}
|
|
if err := db.Model(&entity.IngestionTask{}).Where("id = ?", "task-1").
|
|
Update("status", common.COMPLETED).Error; err != nil {
|
|
t.Fatalf("set COMPLETED: %v", err)
|
|
}
|
|
|
|
svc := NewIngestionTaskService()
|
|
task, err := svc.StartRunning("task-1")
|
|
if err != nil {
|
|
t.Fatalf("StartRunning failed: %v", err)
|
|
}
|
|
if task.Status != common.COMPLETED {
|
|
t.Fatalf("status = %q, want %q (terminal must be preserved)", task.Status, common.COMPLETED)
|
|
}
|
|
|
|
var doc entity.Document
|
|
if err := db.Where("id = ?", "doc-1").First(&doc).Error; err != nil {
|
|
t.Fatalf("reload document: %v", err)
|
|
}
|
|
if doc.Run == nil || *doc.Run != finishedRun {
|
|
t.Fatalf("run = %v, want %q (terminal document must not be resurrected)", doc.Run, finishedRun)
|
|
}
|
|
if doc.ChunkNum != 10 || doc.TokenNum != 100 {
|
|
t.Fatalf("counters changed: chunk_num=%d token_num=%d, want 10/100", doc.ChunkNum, doc.TokenNum)
|
|
}
|
|
}
|
|
|
|
func TestIngestionTaskServiceRequestStopTransitionsCreatedTaskToStopped(t *testing.T) {
|
|
db := setupServiceTestDB(t)
|
|
pushServiceDB(t, db)
|
|
insertTestIngestionTask(t, "task-1", "user-1", "doc-1", "kb-1")
|
|
|
|
svc := NewIngestionTaskService()
|
|
task, err := svc.RequestStop("task-1")
|
|
if err != nil {
|
|
t.Fatalf("RequestStop failed: %v", err)
|
|
}
|
|
if task.Status != common.STOPPED {
|
|
t.Fatalf("status = %q, want %q", task.Status, common.STOPPED)
|
|
}
|
|
}
|
|
|
|
func TestIngestionTaskServiceMarkCompletedRejectsNonRunningTask(t *testing.T) {
|
|
db := setupServiceTestDB(t)
|
|
pushServiceDB(t, db)
|
|
insertTestIngestionTask(t, "task-1", "user-1", "doc-1", "kb-1")
|
|
|
|
svc := NewIngestionTaskService()
|
|
if err := svc.MarkCompleted("task-1"); err == nil {
|
|
t.Fatal("expected MarkCompleted to reject non-running task")
|
|
}
|
|
task, err := dao.NewIngestionTaskDAO().GetByID("task-1")
|
|
if err != nil {
|
|
t.Fatalf("load task: %v", err)
|
|
}
|
|
if task.Status != common.CREATED {
|
|
t.Fatalf("status = %q, want %q", task.Status, common.CREATED)
|
|
}
|
|
}
|
|
|
|
func TestIngestionTaskServiceMarkCompletedUpdatesTaskStatus(t *testing.T) {
|
|
db := setupServiceTestDB(t)
|
|
pushServiceDB(t, db)
|
|
insertTestIngestionTask(t, "task-1", "user-1", "doc-1", "kb-1")
|
|
if err := dao.DB.Model(&entity.IngestionTask{}).Where("id = ?", "task-1").Update("status", common.RUNNING).Error; err != nil {
|
|
t.Fatalf("set running status: %v", err)
|
|
}
|
|
|
|
svc := NewIngestionTaskService()
|
|
if err := svc.MarkCompleted("task-1"); err != nil {
|
|
t.Fatalf("MarkCompleted failed: %v", err)
|
|
}
|
|
task, err := dao.NewIngestionTaskDAO().GetByID("task-1")
|
|
if err != nil {
|
|
t.Fatalf("load task: %v", err)
|
|
}
|
|
if task.Status != common.COMPLETED {
|
|
t.Fatalf("status = %q, want %q", task.Status, common.COMPLETED)
|
|
}
|
|
}
|
|
|
|
func TestIngestionTaskServiceMarkFailedUpdatesTaskStatus(t *testing.T) {
|
|
db := setupServiceTestDB(t)
|
|
pushServiceDB(t, db)
|
|
insertTestIngestionTask(t, "task-1", "user-1", "doc-1", "kb-1")
|
|
if err := dao.DB.Model(&entity.IngestionTask{}).Where("id = ?", "task-1").Update("status", common.RUNNING).Error; err != nil {
|
|
t.Fatalf("set running status: %v", err)
|
|
}
|
|
|
|
svc := NewIngestionTaskService()
|
|
if err := svc.MarkFailed("task-1"); err != nil {
|
|
t.Fatalf("MarkFailed failed: %v", err)
|
|
}
|
|
task, err := dao.NewIngestionTaskDAO().GetByID("task-1")
|
|
if err != nil {
|
|
t.Fatalf("load task: %v", err)
|
|
}
|
|
if task.Status != common.FAILED {
|
|
t.Fatalf("status = %q, want %q", task.Status, common.FAILED)
|
|
}
|
|
}
|
|
|
|
func TestIngestionTaskServiceNewTaskStatusConflictErrorLoadsActualStatus(t *testing.T) {
|
|
db := setupServiceTestDB(t)
|
|
pushServiceDB(t, db)
|
|
insertTestIngestionTask(t, "task-1", "user-1", "doc-1", "kb-1")
|
|
if err := dao.DB.Model(&entity.IngestionTask{}).Where("id = ?", "task-1").Update("status", common.STOPPING).Error; err != nil {
|
|
t.Fatalf("set stopping status: %v", err)
|
|
}
|
|
|
|
svc := NewIngestionTaskService()
|
|
err := svc.newTaskStatusConflictError("task-1", common.CREATED, common.RUNNING)
|
|
var conflictErr *TaskStatusConflictError
|
|
if !errors.As(err, &conflictErr) {
|
|
t.Fatalf("expected TaskStatusConflictError, got %T", err)
|
|
}
|
|
if conflictErr.TaskID != "task-1" || conflictErr.ExpectedFrom != common.CREATED || conflictErr.AttemptedTo != common.RUNNING || conflictErr.ActualCurrent != common.STOPPING {
|
|
t.Fatalf("unexpected conflict error: %+v", conflictErr)
|
|
}
|
|
}
|
|
|
|
func TestIngestionTaskServiceMarkCompletedReturnsTaskIDInTransitionError(t *testing.T) {
|
|
db := setupServiceTestDB(t)
|
|
pushServiceDB(t, db)
|
|
insertTestIngestionTask(t, "task-1", "user-1", "doc-1", "kb-1")
|
|
|
|
svc := NewIngestionTaskService()
|
|
err := svc.MarkCompleted("task-1")
|
|
var transitionErr *InvalidTaskTransitionError
|
|
if !errors.As(err, &transitionErr) {
|
|
t.Fatalf("expected InvalidTaskTransitionError, got %T", err)
|
|
}
|
|
if transitionErr.TaskID != "task-1" || transitionErr.From != common.CREATED || transitionErr.To != common.COMPLETED {
|
|
t.Fatalf("unexpected transition error: %+v", transitionErr)
|
|
}
|
|
}
|
|
|
|
func TestIngestionTaskServiceCreateAndEnqueueRetriesTerminalTask(t *testing.T) {
|
|
db := setupServiceTestDB(t)
|
|
pushServiceDB(t, db)
|
|
publisher := &recordingTaskPublisher{}
|
|
svc := NewIngestionTaskService()
|
|
svc.taskPublisher = publisher
|
|
|
|
testCases := []struct {
|
|
name string
|
|
status string
|
|
}{
|
|
{name: "failed", status: common.FAILED},
|
|
{name: "stopped", status: common.STOPPED},
|
|
}
|
|
|
|
for _, tc := range testCases {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
publisher.subject = ""
|
|
publisher.messages = nil
|
|
if err := dao.DB.Where("id = ?", "task-1").Delete(&entity.IngestionTask{}).Error; err != nil {
|
|
t.Fatalf("clear task: %v", err)
|
|
}
|
|
insertTestIngestionTask(t, "task-1", "user-1", "doc-1", "kb-1")
|
|
if err := dao.DB.Model(&entity.IngestionTask{}).Where("id = ?", "task-1").Update("status", tc.status).Error; err != nil {
|
|
t.Fatalf("set terminal status: %v", err)
|
|
}
|
|
|
|
task, err := svc.CreateAndEnqueue(&entity.IngestionTask{
|
|
DocumentID: "doc-1",
|
|
UserID: "user-1",
|
|
DatasetID: "kb-1",
|
|
Status: common.CREATED,
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("CreateAndEnqueue failed: %v", err)
|
|
}
|
|
if task.ID != "task-1" {
|
|
t.Fatalf("task ID = %q, want task-1", task.ID)
|
|
}
|
|
if task.Status != common.CREATED {
|
|
t.Fatalf("status = %q, want %q", task.Status, common.CREATED)
|
|
}
|
|
if len(publisher.messages) != 1 || publisher.messages[0].TaskID != "task-1" {
|
|
t.Fatalf("unexpected published messages: %+v", publisher.messages)
|
|
}
|
|
reloaded, err := dao.NewIngestionTaskDAO().GetByID("task-1")
|
|
if err != nil {
|
|
t.Fatalf("reload task: %v", err)
|
|
}
|
|
if reloaded.Status != common.CREATED {
|
|
t.Fatalf("reloaded status = %q, want %q", reloaded.Status, common.CREATED)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestIngestionTaskServiceCreateAndEnqueueRejectsActiveExistingTask(t *testing.T) {
|
|
db := setupServiceTestDB(t)
|
|
pushServiceDB(t, db)
|
|
insertTestIngestionTask(t, "task-1", "user-1", "doc-1", "kb-1")
|
|
publisher := &recordingTaskPublisher{}
|
|
svc := NewIngestionTaskService()
|
|
svc.taskPublisher = publisher
|
|
|
|
_, err := svc.CreateAndEnqueue(&entity.IngestionTask{DocumentID: "doc-1", UserID: "user-1", DatasetID: "kb-1", Status: common.CREATED})
|
|
if err == nil {
|
|
t.Fatal("expected CreateAndEnqueue to reject existing created task")
|
|
}
|
|
if len(publisher.messages) != 0 {
|
|
t.Fatalf("expected no published messages, got %+v", publisher.messages)
|
|
}
|
|
}
|
|
|
|
func TestIngestionTaskServiceCreateAndEnqueueRollsBackNewTaskOnPublishFailure(t *testing.T) {
|
|
db := setupServiceTestDB(t)
|
|
pushServiceDB(t, db)
|
|
publisher := &recordingTaskPublisher{err: errors.New("publish failed")}
|
|
svc := NewIngestionTaskService()
|
|
svc.taskPublisher = publisher
|
|
|
|
_, err := svc.CreateAndEnqueue(&entity.IngestionTask{
|
|
DocumentID: "doc-1",
|
|
UserID: "user-1",
|
|
DatasetID: "kb-1",
|
|
Status: common.CREATED,
|
|
})
|
|
if err == nil || err.Error() != "publish failed" {
|
|
t.Fatalf("expected publish failure, got %v", err)
|
|
}
|
|
task, getErr := dao.NewIngestionTaskDAO().GetByDocumentID("doc-1")
|
|
if getErr != nil {
|
|
t.Fatalf("reload task by document id: %v", getErr)
|
|
}
|
|
if task != nil {
|
|
t.Fatalf("expected created task to be deleted after publish failure, got %+v", task)
|
|
}
|
|
}
|
|
|
|
func TestIngestionTaskServiceCreateAndEnqueueRollsBackRetriedTaskOnPublishFailure(t *testing.T) {
|
|
db := setupServiceTestDB(t)
|
|
pushServiceDB(t, db)
|
|
insertTestIngestionTask(t, "task-1", "user-1", "doc-1", "kb-1")
|
|
if err := dao.DB.Model(&entity.IngestionTask{}).Where("id = ?", "task-1").Update("status", common.FAILED).Error; err != nil {
|
|
t.Fatalf("set failed status: %v", err)
|
|
}
|
|
|
|
publisher := &recordingTaskPublisher{err: errors.New("publish failed")}
|
|
svc := NewIngestionTaskService()
|
|
svc.taskPublisher = publisher
|
|
|
|
_, err := svc.CreateAndEnqueue(&entity.IngestionTask{
|
|
DocumentID: "doc-1",
|
|
UserID: "user-1",
|
|
DatasetID: "kb-1",
|
|
Status: common.CREATED,
|
|
})
|
|
if err == nil || err.Error() != "publish failed" {
|
|
t.Fatalf("expected publish failure, got %v", err)
|
|
}
|
|
reloaded, getErr := dao.NewIngestionTaskDAO().GetByID("task-1")
|
|
if getErr != nil {
|
|
t.Fatalf("reload task: %v", getErr)
|
|
}
|
|
if reloaded.Status != common.FAILED {
|
|
t.Fatalf("status = %q, want %q", reloaded.Status, common.FAILED)
|
|
}
|
|
}
|
|
|
|
func TestIngestionTaskServiceRemoveDeletesOwnedTask(t *testing.T) {
|
|
db := setupServiceTestDB(t)
|
|
pushServiceDB(t, db)
|
|
insertTestIngestionTask(t, "task-1", "user-1", "doc-1", "kb-1")
|
|
|
|
userID := "user-1"
|
|
svc := NewIngestionTaskService()
|
|
info, err := svc.Remove("task-1", &userID)
|
|
if err != nil {
|
|
t.Fatalf("Remove failed: %v", err)
|
|
}
|
|
if info == nil || info.TaskID != "task-1" {
|
|
t.Fatalf("unexpected task info: %+v", info)
|
|
}
|
|
if _, err := dao.NewIngestionTaskDAO().GetByID("task-1"); err == nil {
|
|
t.Fatal("task should be removed")
|
|
}
|
|
}
|
|
|
|
func TestIngestionTaskServiceUpdateComponentTotalPersistsDenominator(t *testing.T) {
|
|
db := setupServiceTestDB(t)
|
|
pushServiceDB(t, db)
|
|
insertTestIngestionTask(t, "task-1", "user-1", "doc-1", "kb-1")
|
|
|
|
svc := NewIngestionTaskService()
|
|
if err := svc.UpdateComponentTotal("task-1", 4); err != nil {
|
|
t.Fatalf("UpdateComponentTotal failed: %v", err)
|
|
}
|
|
task, err := dao.NewIngestionTaskDAO().GetByID("task-1")
|
|
if err != nil {
|
|
t.Fatalf("load task: %v", err)
|
|
}
|
|
if task.ComponentTotal != 4 {
|
|
t.Fatalf("component_total = %d, want 4", task.ComponentTotal)
|
|
}
|
|
}
|
|
|
|
func TestIngestionTaskServiceRecordComponentProgressAppendsRow(t *testing.T) {
|
|
db := setupServiceTestDB(t)
|
|
pushServiceDB(t, db)
|
|
insertTestIngestionTask(t, "task-1", "user-1", "doc-1", "kb-1")
|
|
|
|
svc := NewIngestionTaskService()
|
|
if err := svc.RecordComponentProgress("task-1", "Parser", 1, "Parser Done"); err != nil {
|
|
t.Fatalf("RecordComponentProgress failed: %v", err)
|
|
}
|
|
logs, err := dao.NewIngestionTaskLogDAO().ListLogsByTaskID("task-1")
|
|
if err != nil {
|
|
t.Fatalf("list logs: %v", err)
|
|
}
|
|
if len(logs) != 1 {
|
|
t.Fatalf("expected 1 log row, got %d", len(logs))
|
|
}
|
|
row := logs[0]
|
|
if row.Component != "Parser" || row.Phase != 1 || row.Message != "Parser Done" {
|
|
t.Fatalf("unexpected log row: %+v", row)
|
|
}
|
|
if len(row.Checkpoint) != 0 {
|
|
t.Fatalf("component progress row must have empty checkpoint, got %v", row.Checkpoint)
|
|
}
|
|
}
|
|
|
|
func TestIngestionTaskServiceAggregateTaskProgressClassifiesByPhase(t *testing.T) {
|
|
db := setupServiceTestDB(t)
|
|
pushServiceDB(t, db)
|
|
insertTestIngestionTask(t, "task-1", "user-1", "doc-1", "kb-1")
|
|
|
|
svc := NewIngestionTaskService()
|
|
if err := svc.RecordComponentProgress("task-1", "Parser", 1, "Parser Done"); err != nil {
|
|
t.Fatalf("record Parser: %v", err)
|
|
}
|
|
if err := svc.RecordComponentProgress("task-1", "Chunker", 0, "Chunker Started"); err != nil {
|
|
t.Fatalf("record Chunker: %v", err)
|
|
}
|
|
agg, err := svc.AggregateTaskProgress("task-1", 2)
|
|
if err != nil {
|
|
t.Fatalf("AggregateTaskProgress failed: %v", err)
|
|
}
|
|
if agg.Done != 1 || agg.Running != 1 || agg.Failed != 0 {
|
|
t.Fatalf("aggregate = %+v, want Done=1 Running=1 Failed=0", agg)
|
|
}
|
|
if agg.Percent != 50 {
|
|
t.Fatalf("percent = %v, want 50", agg.Percent)
|
|
}
|
|
}
|
|
|
|
func TestIngestionTaskServiceIncrementRunCountInitializesAndBumps(t *testing.T) {
|
|
db := setupServiceTestDB(t)
|
|
pushServiceDB(t, db)
|
|
insertTestIngestionTask(t, "task-1", "user-1", "doc-1", "kb-1")
|
|
|
|
svc := NewIngestionTaskService()
|
|
if err := svc.IncrementRunCount("task-1"); err != nil {
|
|
t.Fatalf("IncrementRunCount (first call) failed: %v", err)
|
|
}
|
|
run, ok := svc.lastRunCount("task-1")
|
|
if !ok || run != 1 {
|
|
t.Fatalf("run_count = %v (ok=%v), want 1", run, ok)
|
|
}
|
|
|
|
// Second call bumps the existing counter to 2.
|
|
if err := svc.IncrementRunCount("task-1"); err != nil {
|
|
t.Fatalf("IncrementRunCount (second call) failed: %v", err)
|
|
}
|
|
run, _ = svc.lastRunCount("task-1")
|
|
if run != 2 {
|
|
t.Fatalf("run_count after second bump = %v, want 2", run)
|
|
}
|
|
}
|
|
|
|
func TestIngestionTaskServiceIncrementRunCountSkippedCorruptedRunCount(t *testing.T) {
|
|
db := setupServiceTestDB(t)
|
|
pushServiceDB(t, db)
|
|
insertTestIngestionTask(t, "task-1", "user-1", "doc-1", "kb-1")
|
|
if err := dao.DB.Create(&entity.IngestionTaskLog{
|
|
TaskID: "task-1",
|
|
Checkpoint: entity.JSONMap{"run_count": "not-a-number"},
|
|
}).Error; err != nil {
|
|
t.Fatalf("insert bad task log: %v", err)
|
|
}
|
|
|
|
svc := NewIngestionTaskService()
|
|
// Corrupted value is skipped; a fresh run_count=1 row is created.
|
|
if err := svc.IncrementRunCount("task-1"); err != nil {
|
|
t.Fatalf("IncrementRunCount should skip corrupted value, got: %v", err)
|
|
}
|
|
run, ok := svc.lastRunCount("task-1")
|
|
if !ok || run != 1 {
|
|
t.Fatalf("run_count = %v (ok=%v), want 1", run, ok)
|
|
}
|
|
}
|
|
|
|
func TestIngestionTaskServiceIncrementRunCountRecoversFromComponentProgressLog(t *testing.T) {
|
|
db := setupServiceTestDB(t)
|
|
pushServiceDB(t, db)
|
|
insertTestIngestionTask(t, "task-1", "user-1", "doc-1", "kb-1")
|
|
|
|
// Simulate a previous run that created some component-progress logs
|
|
// but died before recording a run_count row. The latest log has no run_count.
|
|
svc := NewIngestionTaskService()
|
|
if err := svc.RecordComponentProgress("task-1", "Parser", 1, "Parser Done"); err != nil {
|
|
t.Fatalf("record Parser: %v", err)
|
|
}
|
|
if err := svc.RecordComponentProgress("task-1", "Chunker", 1, "Chunker Done"); err != nil {
|
|
t.Fatalf("record Chunker: %v", err)
|
|
}
|
|
// Verify latest log has empty checkpoint (no run_count).
|
|
latest, err := dao.NewIngestionTaskLogDAO().LatestLogByTaskID("task-1")
|
|
if err != nil {
|
|
t.Fatalf("load latest: %v", err)
|
|
}
|
|
if len(latest.Checkpoint) != 0 {
|
|
t.Fatalf("component-progress row should have empty checkpoint, got %v", latest.Checkpoint)
|
|
}
|
|
|
|
// IncrementRunCount should create a new row with run_count=1.
|
|
if err := svc.IncrementRunCount("task-1"); err != nil {
|
|
t.Fatalf("IncrementRunCount failed: %v", err)
|
|
}
|
|
run, ok := svc.lastRunCount("task-1")
|
|
if !ok || run != 1 {
|
|
t.Fatalf("run_count = %v (ok=%v), want 1", run, ok)
|
|
}
|
|
|
|
// AggregateProgress should still work (run_count row with component=""
|
|
// has phase=0, which doesn't affect counts).
|
|
agg, err := svc.AggregateTaskProgress("task-1", 2)
|
|
if err != nil {
|
|
t.Fatalf("AggregateTaskProgress: %v", err)
|
|
}
|
|
if agg.Done != 2 {
|
|
t.Fatalf("Done = %d, want 2 (run_count row didn't interfere)", agg.Done)
|
|
}
|
|
}
|
|
|
|
func TestIngestionTaskServiceIncrementRunCountAccumulatesAcrossRetries(t *testing.T) {
|
|
db := setupServiceTestDB(t)
|
|
pushServiceDB(t, db)
|
|
insertTestIngestionTask(t, "task-1", "user-1", "doc-1", "kb-1")
|
|
|
|
svc := NewIngestionTaskService()
|
|
|
|
// First attempt: IncrementRunCount creates run_count=1.
|
|
if err := svc.IncrementRunCount("task-1"); err != nil {
|
|
t.Fatalf("first IncrementRunCount: %v", err)
|
|
}
|
|
// Simulate first run: some components progress, then failure.
|
|
if err := svc.RecordComponentProgress("task-1", "Parser", 1, "Parser Done"); err != nil {
|
|
t.Fatalf("record Parser: %v", err)
|
|
}
|
|
|
|
// Second attempt (retry): should find previous run_count=1 and create row with run_count=2.
|
|
if err := svc.IncrementRunCount("task-1"); err != nil {
|
|
t.Fatalf("second IncrementRunCount: %v", err)
|
|
}
|
|
// More progress, then failure.
|
|
if err := svc.RecordComponentProgress("task-1", "Chunker", 1, "Chunker Done"); err != nil {
|
|
t.Fatalf("record Chunker: %v", err)
|
|
}
|
|
|
|
// Third attempt (retry): should find previous run_count=2 and create row with run_count=3.
|
|
if err := svc.IncrementRunCount("task-1"); err != nil {
|
|
t.Fatalf("third IncrementRunCount: %v", err)
|
|
}
|
|
|
|
run, ok := svc.lastRunCount("task-1")
|
|
if !ok || run != 3 {
|
|
t.Fatalf("run_count = %v (ok=%v), want 3", run, ok)
|
|
}
|
|
|
|
// ListAllForAdmin should still pick up the correct run_count.
|
|
status := "1"
|
|
if err := dao.DB.Create(&entity.User{
|
|
ID: "user-1",
|
|
Email: "user-1@test.com",
|
|
Nickname: "user-1",
|
|
IsAuthenticated: "1",
|
|
IsActive: "1",
|
|
IsAnonymous: "0",
|
|
Status: &status,
|
|
}).Error; err != nil {
|
|
t.Fatalf("insert user: %v", err)
|
|
}
|
|
adminTasks, err := svc.ListAllForAdmin()
|
|
if err != nil {
|
|
t.Fatalf("ListAllForAdmin: %v", err)
|
|
}
|
|
if len(adminTasks) != 1 || adminTasks[0]["id"] != "task-1" {
|
|
t.Fatalf("ListAllForAdmin = %+v, want single task task-1", adminTasks)
|
|
}
|
|
if adminTasks[0]["run_count"] != 3 {
|
|
t.Fatalf("ListAllForAdmin run_count = %v, want 3", adminTasks[0]["run_count"])
|
|
}
|
|
}
|
|
|
|
func TestIngestionTaskServiceMarkStoppedTransitionsStoppingTask(t *testing.T) {
|
|
db := setupServiceTestDB(t)
|
|
pushServiceDB(t, db)
|
|
insertTestIngestionTask(t, "task-1", "user-1", "doc-1", "kb-1")
|
|
// Override to STOPPING.
|
|
if err := dao.DB.Model(&entity.IngestionTask{}).Where("id = ?", "task-1").
|
|
Update("status", common.STOPPING).Error; err != nil {
|
|
t.Fatalf("set task STOPPING: %v", err)
|
|
}
|
|
|
|
svc := NewIngestionTaskService()
|
|
if err := svc.MarkStopped("task-1"); err != nil {
|
|
t.Fatalf("MarkStopped failed: %v", err)
|
|
}
|
|
|
|
task, err := dao.NewIngestionTaskDAO().GetByID("task-1")
|
|
if err != nil {
|
|
t.Fatalf("load task: %v", err)
|
|
}
|
|
if task.Status != common.STOPPED {
|
|
t.Fatalf("task status = %s, want STOPPED", task.Status)
|
|
}
|
|
}
|
|
|
|
func TestIngestionTaskServiceMarkStoppedIdempotentOnAlreadyStopped(t *testing.T) {
|
|
db := setupServiceTestDB(t)
|
|
pushServiceDB(t, db)
|
|
insertTestIngestionTask(t, "task-1", "user-1", "doc-1", "kb-1")
|
|
if err := dao.DB.Model(&entity.IngestionTask{}).Where("id = ?", "task-1").
|
|
Update("status", common.STOPPED).Error; err != nil {
|
|
t.Fatalf("set task STOPPED: %v", err)
|
|
}
|
|
|
|
svc := NewIngestionTaskService()
|
|
if err := svc.MarkStopped("task-1"); err != nil {
|
|
t.Fatalf("MarkStopped on already STOPPED task should be idempotent, got: %v", err)
|
|
}
|
|
}
|
|
|
|
func TestIngestionTaskServiceMarkFailedIdempotentOnAlreadyTerminal(t *testing.T) {
|
|
db := setupServiceTestDB(t)
|
|
pushServiceDB(t, db)
|
|
insertTestIngestionTask(t, "task-1", "user-1", "doc-1", "kb-1")
|
|
if err := dao.DB.Model(&entity.IngestionTask{}).Where("id = ?", "task-1").
|
|
Update("status", common.COMPLETED).Error; err != nil {
|
|
t.Fatalf("set task COMPLETED: %v", err)
|
|
}
|
|
|
|
svc := NewIngestionTaskService()
|
|
if err := svc.MarkFailed("task-1"); err != nil {
|
|
t.Fatalf("MarkFailed on already COMPLETED task should be idempotent, got: %v", err)
|
|
}
|
|
}
|
|
|
|
func TestIngestionTaskServiceMarkCompletedIdempotentOnAlreadyTerminal(t *testing.T) {
|
|
db := setupServiceTestDB(t)
|
|
pushServiceDB(t, db)
|
|
insertTestIngestionTask(t, "task-1", "user-1", "doc-1", "kb-1")
|
|
if err := dao.DB.Model(&entity.IngestionTask{}).Where("id = ?", "task-1").
|
|
Update("status", common.FAILED).Error; err != nil {
|
|
t.Fatalf("set task FAILED: %v", err)
|
|
}
|
|
|
|
svc := NewIngestionTaskService()
|
|
if err := svc.MarkCompleted("task-1"); err != nil {
|
|
t.Fatalf("MarkCompleted on already FAILED task should be idempotent, got: %v", err)
|
|
}
|
|
}
|
|
|
|
func TestDocumentServiceUpdateRunProgressMirrorsFields(t *testing.T) {
|
|
db := setupServiceTestDB(t)
|
|
pushServiceDB(t, db)
|
|
insertTestDoc(t, "doc-1", "kb-1", 0, 0)
|
|
|
|
svc := testDocumentService(t)
|
|
if err := svc.UpdateRunProgress("doc-1", 0.5, "1", "halfway"); err != nil {
|
|
t.Fatalf("UpdateRunProgress failed: %v", err)
|
|
}
|
|
doc, err := dao.NewDocumentDAO().GetByID("doc-1")
|
|
if err != nil {
|
|
t.Fatalf("load document: %v", err)
|
|
}
|
|
if doc.Progress != 0.5 {
|
|
t.Fatalf("progress = %v, want 0.5", doc.Progress)
|
|
}
|
|
if doc.Run == nil || *doc.Run != "1" {
|
|
t.Fatalf("run = %v, want 1", doc.Run)
|
|
}
|
|
if doc.ProgressMsg == nil || *doc.ProgressMsg != "halfway" {
|
|
t.Fatalf("progress_msg = %v, want halfway", doc.ProgressMsg)
|
|
}
|
|
}
|