cmsp/internal/task/task_test.go

package task

import (
	"context"
	"errors"
	"reflect"
	"sync"
	"sync/atomic"
	"testing"
	"time"
)

func Test正常跑完N个商品(t *testing.T) {
	var mu sync.Mutex
	var doneValues []int
	runner := newTestRunner(Options{
		OnProgress: func(progress Progress) {
			mu.Lock()
			doneValues = append(doneValues, progress.Done)
			mu.Unlock()
		},
	})
	if err := runner.Start(context.Background(), []string{"1", "2", "3"}); err != nil {
		t.Fatalf("启动任务失败：%v", err)
	}
	progress := waitForState(t, runner, StateCompleted)
	if progress.Total != 3 || progress.Done != 3 {
		t.Fatalf("完成进度错误：%+v", progress)
	}
	mu.Lock()
	defer mu.Unlock()
	if !reflect.DeepEqual(doneValues, []int{1, 2, 3}) {
		t.Fatalf("进度没有从 0 依次推进到 N：%v", doneValues)
	}
}

func Test中途停止后不再启动剩余商品(t *testing.T) {
	started := make(chan string, 3)
	release := make(chan struct{})
	runner := newTestRunner(Options{
		Prepare: func(_ context.Context, product Product) (Work, error) {
			started <- product.ID
			if product.ID == "1" {
				<-release
			}
			return Work{}, nil
		},
	})
	if err := runner.Start(context.Background(), []string{"1", "2", "3"}); err != nil {
		t.Fatalf("启动任务失败：%v", err)
	}
	if got := <-started; got != "1" {
		t.Fatalf("首个商品错误：%s", got)
	}
	runner.Stop()
	close(release)
	progress := waitForState(t, runner, StateStopped)
	if progress.Done != 1 {
		t.Fatalf("当前商品应处理完再停止：%+v", progress)
	}
	select {
	case id := <-started:
		t.Fatalf("停止后仍启动了商品：%s", id)
	default:
	}
}

func Test登录失效立即触发全局停止门(t *testing.T) {
	var processed []string
	var mu sync.Mutex
	runner := newTestRunner(Options{
		Prepare: func(_ context.Context, product Product) (Work, error) {
			mu.Lock()
			processed = append(processed, product.ID)
			mu.Unlock()
			if product.ID == "2" {
				return Work{}, errors.Join(ErrLoginRequired, errors.New("深度检查不通过"))
			}
			return Work{}, nil
		},
	})
	if err := runner.Start(context.Background(), []string{"1", "2", "3", "4"}); err != nil {
		t.Fatalf("启动任务失败：%v", err)
	}
	progress := waitForState(t, runner, StateLoginRequired)
	mu.Lock()
	defer mu.Unlock()
	if !reflect.DeepEqual(processed, []string{"1", "2"}) {
		t.Fatalf("登录失效后仍处理了后续商品：%v", processed)
	}
	if progress.Done != 1 || progress.Failed != 0 {
		t.Fatalf("登录失效商品不能记为单商品失败：%+v", progress)
	}
}

func Test疑似风控触发停止门且后续商品不处理(t *testing.T) {
	var processed []string
	runner := newTestRunner(Options{Prepare: func(_ context.Context, product Product) (Work, error) {
		processed = append(processed, product.ID)
		if product.ID == "2" {
			return Work{}, errors.Join(ErrLoginRequired, ErrRiskSuspected, RiskSuspectedError{Consecutive: 8})
		}
		return Work{}, nil
	}})
	if err := runner.Start(context.Background(), []string{"1", "2", "3"}); err != nil {
		t.Fatalf("启动任务失败：%v", err)
	}
	progress := waitForState(t, runner, StateLoginRequired)
	if !reflect.DeepEqual(processed, []string{"1", "2"}) {
		t.Fatalf("疑似风控后仍处理后续商品：%v", processed)
	}
	if progress.StopReason != StopReasonRiskSuspected {
		t.Fatalf("停止原因应为疑似风控，实际：%+v", progress)
	}
	if progress.RiskEmptyCount != 8 {
		t.Fatalf("连续空结果数应传给前端，实际：%+v", progress)
	}
}

func Test空结果中途发现视频会清零(t *testing.T) {
	guard := NewEmptyRiskGuard(3)
	if guard.Record(0) || guard.Record(0) {
		t.Fatal("阈值前不应触发")
	}
	if guard.Record(1) || guard.Consecutive() != 0 {
		t.Fatalf("发现视频应清零，实际连续数 %d", guard.Consecutive())
	}
	if guard.Record(0) || guard.Record(0) || !guard.Record(0) {
		t.Fatal("清零后应重新累计到阈值")
	}
}

func Test单商品失败后继续处理后续商品(t *testing.T) {
	var processed []string
	var mu sync.Mutex
	runner := newTestRunner(Options{
		Prepare: func(_ context.Context, product Product) (Work, error) {
			mu.Lock()
			processed = append(processed, product.ID)
			mu.Unlock()
			if product.ID == "2" {
				return Work{}, errors.New("图搜失败")
			}
			return Work{}, nil
		},
	})
	if err := runner.Start(context.Background(), []string{"1", "2", "3"}); err != nil {
		t.Fatalf("启动任务失败：%v", err)
	}
	progress := waitForState(t, runner, StateCompleted)
	mu.Lock()
	defer mu.Unlock()
	if !reflect.DeepEqual(processed, []string{"1", "2", "3"}) {
		t.Fatalf("单商品失败后没有继续：%v", processed)
	}
	if progress.Failed != 1 || progress.Done != 3 {
		t.Fatalf("失败计数错误：%+v", progress)
	}
}

func Test已有任务运行时拒绝再次启动(t *testing.T) {
	release := make(chan struct{})
	started := make(chan struct{})
	runner := newTestRunner(Options{
		Prepare: func(context.Context, Product) (Work, error) {
			close(started)
			<-release
			return Work{}, nil
		},
	})
	if err := runner.Start(context.Background(), []string{"1"}); err != nil {
		t.Fatalf("首次启动失败：%v", err)
	}
	<-started
	if err := runner.Start(context.Background(), []string{"2"}); err == nil || err.Error() != "已有任务在运行" {
		t.Fatalf("重复启动应返回指定中文错误，实际：%v", err)
	}
	close(release)
	waitForState(t, runner, StateCompleted)
}

func Test已完成和无同款商品不进入队列(t *testing.T) {
	statuses := map[string]Product{
		"done": {ID: "done", DownloadStatus: "done"},
		"none": {ID: "none", VideoStatus: "none"},
		"todo": {ID: "todo"},
	}
	var processed []string
	var mu sync.Mutex
	runner := newTestRunner(Options{
		Load: func(_ context.Context, id string) (Product, error) { return statuses[id], nil },
		Prepare: func(_ context.Context, product Product) (Work, error) {
			mu.Lock()
			processed = append(processed, product.ID)
			mu.Unlock()
			return Work{}, nil
		},
	})
	if err := runner.Start(context.Background(), []string{"done", "none", "todo"}); err != nil {
		t.Fatalf("启动任务失败：%v", err)
	}
	progress := waitForState(t, runner, StateCompleted)
	mu.Lock()
	defer mu.Unlock()
	if !reflect.DeepEqual(processed, []string{"todo"}) {
		t.Fatalf("断点跳过逻辑错误：%v", processed)
	}
	if progress.Total != 1 || progress.Done != 1 {
		t.Fatalf("跳过后的总数错误：%+v", progress)
	}
}

func Test下载并发数不超过配置值(t *testing.T) {
	var active int32
	var maximum int32
	jobs := make([]DownloadFunc, 8)
	for index := range jobs {
		jobs[index] = func(context.Context) error {
			current := atomic.AddInt32(&active, 1)
			for {
				old := atomic.LoadInt32(&maximum)
				if current <= old || atomic.CompareAndSwapInt32(&maximum, old, current) {
					break
				}
			}
			time.Sleep(15 * time.Millisecond)
			atomic.AddInt32(&active, -1)
			return nil
		}
	}
	runner := newTestRunner(Options{
		Concurrency: 2,
		Prepare: func(context.Context, Product) (Work, error) {
			return Work{Downloads: jobs}, nil
		},
	})
	if err := runner.Start(context.Background(), []string{"1"}); err != nil {
		t.Fatalf("启动任务失败：%v", err)
	}
	progress := waitForState(t, runner, StateCompleted)
	if maximum != 2 {
		t.Fatalf("下载同时运行数应为 2，实际为 %d", maximum)
	}
	if progress.Downloaded != len(jobs) {
		t.Fatalf("下载计数错误：%+v", progress)
	}
}

func Test每处理完一个商品调用一次进度回调(t *testing.T) {
	var calls int32
	runner := newTestRunner(Options{
		OnProgress: func(Progress) { atomic.AddInt32(&calls, 1) },
	})
	if err := runner.Start(context.Background(), []string{"1", "2", "3", "4"}); err != nil {
		t.Fatalf("启动任务失败：%v", err)
	}
	waitForState(t, runner, StateCompleted)
	if calls != 4 {
		t.Fatalf("进度回调次数应等于处理商品数，实际为 %d", calls)
	}
}

func newTestRunner(override Options) *Runner {
	options := Options{
		Concurrency: 1,
		Load: func(_ context.Context, id string) (Product, error) {
			return Product{ID: id, Name: "商品" + id}, nil
		},
		Prepare: func(context.Context, Product) (Work, error) { return Work{}, nil },
	}
	if override.Load != nil {
		options.Load = override.Load
	}
	if override.Prepare != nil {
		options.Prepare = override.Prepare
	}
	if override.Concurrency > 0 {
		options.Concurrency = override.Concurrency
	}
	options.OnProgress = override.OnProgress
	options.OnFinished = override.OnFinished
	return NewRunner(options)
}

func waitForState(t *testing.T, runner *Runner, state string) Progress {
	t.Helper()
	deadline := time.Now().Add(2 * time.Second)
	for time.Now().Before(deadline) {
		progress := runner.Snapshot()
		if progress.State == state {
			return progress
		}
		time.Sleep(time.Millisecond)
	}
	t.Fatalf("等待任务状态 %s 超时，当前：%+v", state, runner.Snapshot())
	return Progress{}
}