309 lines
7.9 KiB
Go
309 lines
7.9 KiB
Go
// 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)
|
|
}
|
|
|
|
func (m *MockTask) GetDescription() string {
|
|
return "Mock Task"
|
|
}
|
|
|
|
// --- 健壮等待的辅助函数 ---
|
|
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:
|
|
// 等待成功
|
|
case <-time.After(timeout):
|
|
t.Fatal("等待任务完成超时")
|
|
}
|
|
}
|
|
|
|
// --- 任务队列测试 (无需更改) ---
|
|
|
|
func TestNewQueue(t *testing.T) {
|
|
tq := task.NewQueue(testLogger)
|
|
assert.NotNil(t, tq, "新创建的任务队列不应为 nil")
|
|
assert.Equal(t, 0, tq.GetTaskCount(), "新创建的任务队列应为空")
|
|
}
|
|
|
|
func TestQueue_AddTask(t *testing.T) {
|
|
tq := task.NewQueue(testLogger)
|
|
mockTask := &MockTask{id: "task1", priority: 1}
|
|
|
|
tq.AddTask(mockTask)
|
|
assert.Equal(t, 1, tq.GetTaskCount(), "添加任务后，队列中的任务数应为 1")
|
|
}
|
|
|
|
// ... (其他任务队列测试保持不变)
|
|
func TestQueue_GetNextTask(t *testing.T) {
|
|
t.Run("从空队列获取任务", func(t *testing.T) {
|
|
tq := task.NewQueue(testLogger)
|
|
nextTask := tq.GetNextTask()
|
|
assert.Nil(t, nextTask, "从空队列中获取任务应返回 nil")
|
|
})
|
|
|
|
t.Run("按优先级获取任务", func(t *testing.T) {
|
|
tq := task.NewQueue(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 TestQueue_Concurrency(t *testing.T) {
|
|
tq := task.NewQueue(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(), "并发获取所有任务后，队列应为空")
|
|
}
|
|
|
|
// --- 执行器测试 (为可靠性重构) ---
|
|
|
|
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("等待 task1 启动超时")
|
|
}
|
|
|
|
executor.Stop()
|
|
|
|
assert.Equal(t, int32(1), task1.ExecutedCount(), "task1 应该在停止前开始执行")
|
|
assert.Equal(t, int32(0), task2.ExecutedCount(), "task2 不应该被执行，因为执行器已停止")
|
|
}
|