// Package task_test 包含对 task 包的单元测试 package task_test import ( "errors" "fmt" "sync" "sync/atomic" "testing" "time" "git.huangwc.com/pig/pig-farm-controller/internal/infra/config" "git.huangwc.com/pig/pig-farm-controller/internal/infra/logs" "git.huangwc.com/pig/pig-farm-controller/internal/infra/task" "github.com/stretchr/testify/assert" ) // testLogger 是一个用于所有测试用例的静默 logger 实例。 var testLogger *logs.Logger func init() { // 使用 "fatal" 级别来创建一个在测试期间不会产生任何输出的 logger。 // 这避免了在运行 `go test` 时被日志淹没。 cfg := config.LogConfig{Level: "fatal"} testLogger = logs.NewLogger(cfg) } // MockTask 用于测试的模拟任务 type MockTask struct { id string priority int isDone bool execute func() error executed int32 // 使用原子操作来跟踪执行次数 } // Execute 实现了 Task 接口,并确保每次调用都增加执行计数 func (m *MockTask) Execute() error { atomic.AddInt32(&m.executed, 1) if m.execute != nil { return m.execute() } return nil } func (m *MockTask) GetID() string { return m.id } func (m *MockTask) GetPriority() int { return m.priority } func (m *MockTask) IsDone() bool { return m.isDone } // ExecutedCount 返回任务被执行的次数 func (m *MockTask) ExecutedCount() int32 { return atomic.LoadInt32(&m.executed) } // --- Helper function for robust waiting --- func waitForWaitGroup(t *testing.T, wg *sync.WaitGroup, timeout time.Duration) { waitChan := make(chan struct{}) go func() { defer close(waitChan) wg.Wait() }() select { case <-waitChan: // Wait succeeded case <-time.After(timeout): t.Fatal("timed out waiting for tasks to complete") } } // --- TaskQueue Tests (No changes needed) --- func TestNewTaskQueue(t *testing.T) { tq := task.NewTaskQueue(testLogger) assert.NotNil(t, tq, "新创建的任务队列不应为 nil") assert.Equal(t, 0, tq.GetTaskCount(), "新创建的任务队列应为空") } func TestTaskQueue_AddTask(t *testing.T) { tq := task.NewTaskQueue(testLogger) mockTask := &MockTask{id: "task1", priority: 1} tq.AddTask(mockTask) assert.Equal(t, 1, tq.GetTaskCount(), "添加任务后,队列中的任务数应为 1") } // ... (other TaskQueue tests remain the same) func TestTaskQueue_GetNextTask(t *testing.T) { t.Run("从空队列获取任务", func(t *testing.T) { tq := task.NewTaskQueue(testLogger) nextTask := tq.GetNextTask() assert.Nil(t, nextTask, "从空队列中获取任务应返回 nil") }) t.Run("按优先级获取任务", func(t *testing.T) { tq := task.NewTaskQueue(testLogger) task1 := &MockTask{id: "task1", priority: 10} task2 := &MockTask{id: "task2", priority: 1} // 优先级更高 task3 := &MockTask{id: "task3", priority: 5} tq.AddTask(task1) tq.AddTask(task2) tq.AddTask(task3) assert.Equal(t, 3, tq.GetTaskCount(), "添加三个任务后,队列中的任务数应为 3") nextTask := tq.GetNextTask() assert.NotNil(t, nextTask) assert.Equal(t, "task2", nextTask.GetID(), "应首先获取优先级最高的任务 (task2)") nextTask = tq.GetNextTask() assert.NotNil(t, nextTask) assert.Equal(t, "task3", nextTask.GetID(), "应获取下一个优先级最高的任务 (task3)") nextTask = tq.GetNextTask() assert.NotNil(t, nextTask) assert.Equal(t, "task1", nextTask.GetID(), "应最后获取优先级最低的任务 (task1)") assert.Equal(t, 0, tq.GetTaskCount(), "获取所有任务后,队列应为空") }) } func TestTaskQueue_Concurrency(t *testing.T) { tq := task.NewTaskQueue(testLogger) var wg sync.WaitGroup taskCount := 100 wg.Add(taskCount) for i := 0; i < taskCount; i++ { go func(i int) { defer wg.Done() tq.AddTask(&MockTask{id: fmt.Sprintf("task-%d", i), priority: i}) }(i) } wg.Wait() assert.Equal(t, taskCount, tq.GetTaskCount(), "并发添加任务后,队列中的任务数应为 %d", taskCount) wg.Add(taskCount) for i := 0; i < taskCount; i++ { go func() { defer wg.Done() task := tq.GetNextTask() assert.NotNil(t, task) }() } wg.Wait() assert.Equal(t, 0, tq.GetTaskCount(), "并发获取所有任务后,队列应为空") } // --- Executor Tests (Refactored for reliability) --- func TestNewExecutor(t *testing.T) { executor := task.NewExecutor(5, testLogger) assert.NotNil(t, executor, "新创建的执行器不应为 nil") } func TestExecutor_StartStop(t *testing.T) { executor := task.NewExecutor(2, testLogger) executor.Start() // 确保立即停止不会导致死锁或竞争条件。 executor.Stop() } // TestExecutor_SubmitAndExecuteTask 测试提交并执行单个任务 (已重构,更可靠) func TestExecutor_SubmitAndExecuteTask(t *testing.T) { var wg sync.WaitGroup wg.Add(1) executor := task.NewExecutor(1, testLogger) mockTask := &MockTask{ id: "task1", priority: 1, execute: func() error { wg.Done() // 任务完成时通知 WaitGroup return nil }, } executor.Start() executor.SubmitTask(mockTask) // 等待任务完成,设置一个合理的超时时间 waitForWaitGroup(t, &wg, 2*time.Second) executor.Stop() assert.Equal(t, int32(1), mockTask.ExecutedCount(), "任务应该已被执行") } // TestExecutor_ExecuteMultipleTasks 测试执行多个任务 (已重构,更可靠) func TestExecutor_ExecuteMultipleTasks(t *testing.T) { taskCount := 10 var wg sync.WaitGroup wg.Add(taskCount) executor := task.NewExecutor(3, testLogger) mockTasks := make([]*MockTask, taskCount) for i := 0; i < taskCount; i++ { mockTasks[i] = &MockTask{ id: fmt.Sprintf("task-%d", i), priority: i, execute: func() error { wg.Done() // 每个任务完成时都通知 WaitGroup return nil }, } } executor.Start() for _, task := range mockTasks { executor.SubmitTask(task) } // 等待所有任务完成 waitForWaitGroup(t, &wg, 2*time.Second) executor.Stop() var totalExecuted int32 for _, task := range mockTasks { totalExecuted += task.ExecutedCount() } assert.Equal(t, int32(taskCount), totalExecuted, "所有提交的任务都应该被执行") } // TestExecutor_TaskExecutionError 测试任务执行失败的场景 (已重构,更可靠) func TestExecutor_TaskExecutionError(t *testing.T) { var wg sync.WaitGroup wg.Add(2) // 我们期望两个任务都被执行 executor := task.NewExecutor(1, testLogger) errorTask := &MockTask{ id: "errorTask", priority: 1, execute: func() error { wg.Done() return errors.New("执行失败") }, } successTask := &MockTask{ id: "successTask", priority: 2, // 后执行 execute: func() error { wg.Done() return nil }, } executor.Start() executor.SubmitTask(errorTask) executor.SubmitTask(successTask) waitForWaitGroup(t, &wg, 2*time.Second) executor.Stop() assert.Equal(t, int32(1), errorTask.ExecutedCount(), "失败的任务应该被执行一次") assert.Equal(t, int32(1), successTask.ExecutedCount(), "成功的任务也应该被执行") } // TestExecutor_StopWithPendingTasks 测试停止执行器时仍有待处理任务 (已重构,更可靠) func TestExecutor_StopWithPendingTasks(t *testing.T) { executor := task.NewExecutor(1, testLogger) task1Started := make(chan struct{}) task1 := &MockTask{ id: "task1", priority: 1, execute: func() error { close(task1Started) // 发送信号,通知测试 task1 已开始执行 time.Sleep(200 * time.Millisecond) // 模拟耗时操作 return nil }, } task2 := &MockTask{id: "task2", priority: 2} executor.Start() executor.SubmitTask(task1) executor.SubmitTask(task2) // 等待 task1 开始执行的信号,而不是依赖不确定的 sleep select { case <-task1Started: // task1 已开始,可以安全地停止执行器了 case <-time.After(1 * time.Second): t.Fatal("timed out waiting for task1 to start") } executor.Stop() assert.Equal(t, int32(1), task1.ExecutedCount(), "task1 应该在停止前开始执行") assert.Equal(t, int32(0), task2.ExecutedCount(), "task2 不应该被执行,因为执行器已停止") }