From 6ad2cf2003b6eaf1cbe5ad67788c18c6d6740ce4 Mon Sep 17 00:00:00 2001
From: Noah Hsu <i@nn.ci>
Date: Fri, 17 Jun 2022 22:09:20 +0800
Subject: [PATCH] test: add task manager test

---
 pkg/task/manager.go   | 19 +++++++++++
 pkg/task/task_test.go | 79 +++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 98 insertions(+)
 create mode 100644 pkg/task/task_test.go

diff --git a/pkg/task/manager.go b/pkg/task/manager.go
index 4d3874b8..4dfd65ca 100644
--- a/pkg/task/manager.go
+++ b/pkg/task/manager.go
@@ -1,6 +1,7 @@
 package task
 
 import (
+	"github.com/pkg/errors"
 	"sync/atomic"
 
 	"github.com/alist-org/alist/v3/pkg/generic_sync"
@@ -32,6 +33,24 @@ func (tm *Manager) Get(id uint64) (*Task, bool) {
 	return tm.tasks.Load(id)
 }
 
+func (tm *Manager) Retry(id uint64) error {
+	t, ok := tm.Get(id)
+	if !ok {
+		return errors.New("task not found")
+	}
+	t.Retry()
+	return nil
+}
+
+func (tm *Manager) Cancel(id uint64) error {
+	t, ok := tm.Get(id)
+	if !ok {
+		return errors.New("task not found")
+	}
+	t.Cancel()
+	return nil
+}
+
 func (tm *Manager) Remove(id uint64) {
 	tm.tasks.Delete(id)
 }
diff --git a/pkg/task/task_test.go b/pkg/task/task_test.go
new file mode 100644
index 00000000..352781aa
--- /dev/null
+++ b/pkg/task/task_test.go
@@ -0,0 +1,79 @@
+package task
+
+import (
+	"github.com/alist-org/alist/v3/pkg/utils"
+	"github.com/pkg/errors"
+	"testing"
+	"time"
+)
+
+func TestTask_Manager(t *testing.T) {
+	tm := NewTaskManager()
+	id := tm.Add("test", func(task *Task) error {
+		time.Sleep(time.Millisecond * 500)
+		return nil
+	})
+	task, ok := tm.Get(id)
+	if !ok {
+		t.Fatal("task not found")
+	}
+	time.Sleep(time.Millisecond * 100)
+	if task.Status != RUNNING {
+		t.Errorf("task status not running: %s", task.Status)
+	}
+	time.Sleep(time.Second)
+	if task.Status != FINISHED {
+		t.Errorf("task status not finished: %s", task.Status)
+	}
+}
+
+func TestTask_Cancel(t *testing.T) {
+	tm := NewTaskManager()
+	id := tm.Add("test", func(task *Task) error {
+		for {
+			if utils.IsCanceled(task.Ctx) {
+				return nil
+			} else {
+				t.Logf("task is running")
+			}
+		}
+	})
+	task, ok := tm.Get(id)
+	if !ok {
+		t.Fatal("task not found")
+	}
+	time.Sleep(time.Microsecond * 50)
+	task.Cancel()
+	time.Sleep(time.Millisecond)
+	if task.Status != CANCELED {
+		t.Errorf("task status not canceled: %s", task.Status)
+	}
+}
+
+func TestTask_Retry(t *testing.T) {
+	tm := NewTaskManager()
+	num := 0
+	id := tm.Add("test", func(task *Task) error {
+		num++
+		if num&1 == 1 {
+			return errors.New("test error")
+		}
+		return nil
+	})
+	task, ok := tm.Get(id)
+	if !ok {
+		t.Fatal("task not found")
+	}
+	time.Sleep(time.Millisecond)
+	if task.Error == nil {
+		t.Error(task.Status)
+		t.Fatal("task error is nil, but expected error")
+	} else {
+		t.Logf("task error: %s", task.Error)
+	}
+	task.Retry()
+	time.Sleep(time.Millisecond)
+	if task.Error != nil {
+		t.Errorf("task error: %+v, but expected nil", task.Error)
+	}
+}