kra-new/internal/worker/task_executor_test.go

package worker

import (
	"context"
	"encoding/json"
	"errors"
	"net"
	"strings"
	"testing"
	"time"

	taskbiz "kra/internal/biz/task"
	platformtask "kra/pkg/task"
)

func TestPrivateIP(t *testing.T) {
	for _, value := range []string{"127.0.0.1", "10.0.0.1", "172.16.0.1", "192.168.1.1", "169.254.1.1", "::1"} {
		if !privateIP(net.ParseIP(value)) {
			t.Fatalf("expected %s to be private", value)
		}
	}
	if privateIP(net.ParseIP("8.8.8.8")) {
		t.Fatal("public IP was classified as private")
	}
}

func TestMethodExecutionPersistsProgressAndRetryOutcome(t *testing.T) {
	registry := platformtask.NewRegistry()
	attempts := 0
	registry.Register(platformtask.Method{Name: "report", Run: func(ctx context.Context, _ json.RawMessage) error {
		attempts++
		if attempts == 1 {
			return errors.New("temporary storage failure")
		}
		platformtask.Logf(ctx, "日期=2026-11-05；已提交 12 行")
		return nil
	}})
	repo := &workerTaskRepo{logs: make(chan *taskbiz.TimedTaskLog, 1)}
	executor := NewTaskExecutorWithRegistry(taskbiz.NewTaskUsecaseWithRegistry(repo, registry), registry)
	log := executor.Run(context.Background(), &taskbiz.TimedTask{ID: 1, Name: "日报", ExecutorType: taskbiz.TaskExecutorMethod, MethodName: "report", Params: []byte(`{"_execution":{"maxRetryTimes":1,"retryIntervalSeconds":0}}`)}, "manual")
	stored := <-repo.logs
	if log.Status != "success" || stored.Output != log.Output {
		t.Fatalf("result was not persisted: %+v", log)
	}
	for _, detail := range []string{"方法：report", "第 1 次执行失败", "第 2 次执行", "已提交 12 行", "执行结束：success"} {
		if !strings.Contains(log.Output, detail) {
			t.Fatalf("missing %q in %s", detail, log.Output)
		}
	}
}

func TestMethodTimeoutKeepsPartialProgressWithoutClaimingCompletion(t *testing.T) {
	registry := platformtask.NewRegistry()
	release := make(chan struct{})
	finished := make(chan struct{})
	registry.Register(platformtask.Method{Name: "slow", Run: func(ctx context.Context, _ json.RawMessage) error {
		defer close(finished)
		platformtask.Logf(ctx, "已处理第一批")
		<-release
		platformtask.Logf(ctx, "迟到的结果")
		return nil
	}})
	ctx, cancel := context.WithTimeout(context.Background(), 25*time.Millisecond)
	defer cancel()
	log := (&TaskExecutor{methods: registry}).Run(ctx, &taskbiz.TimedTask{ID: 1, ExecutorType: taskbiz.TaskExecutorMethod, MethodName: "slow"}, "manual")
	close(release)
	<-finished
	if log.Status != "timeout" || !strings.Contains(log.Output, "已处理第一批") || strings.Contains(log.Output, "迟到的结果") {
		t.Fatalf("timeout did not preserve the execution snapshot: %+v", log)
	}
}

func TestRunMethodReportsParentDeadlineAsTimeout(t *testing.T) {
	const methodName = "worker-test-parent-deadline"
	methods := platformtask.NewRegistry()
	methods.Register(platformtask.Method{Name: methodName, Description: "test", Run: func(ctx context.Context, _ json.RawMessage) error {
		<-ctx.Done()
		return ctx.Err()
	}})
	executor := &TaskExecutor{methods: methods}
	ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond)
	defer cancel()
	if err := executor.runMethod(ctx, &taskbiz.TimedTask{MethodName: methodName}); !errors.Is(err, errTaskTimeout) {
		t.Fatalf("runMethod() error = %v, want errTaskTimeout", err)
	}
}

func TestRunMethodConvertsPanicToError(t *testing.T) {
	const methodName = "worker-test-panic"
	methods := platformtask.NewRegistry()
	methods.Register(platformtask.Method{Name: methodName, Description: "test", Run: func(context.Context, json.RawMessage) error {
		panic("boom")
	}})
	executor := &TaskExecutor{methods: methods}
	err := executor.runMethod(context.Background(), &taskbiz.TimedTask{MethodName: methodName})
	if err == nil || !strings.Contains(err.Error(), "panic: boom") {
		t.Fatalf("runMethod() error = %v, want recovered panic", err)
	}
}

func TestRunMethodUsesImportedDeadlineAndRetries(t *testing.T) {
	registry := platformtask.NewRegistry()
	attempts := 0
	registry.Register(platformtask.Method{Name: "migrated-retry", Run: func(ctx context.Context, _ json.RawMessage) error {
		attempts++
		deadline, ok := ctx.Deadline()
		if !ok || time.Until(deadline) < 1590*time.Second || time.Until(deadline) > 1600*time.Second {
			return errors.New("迁移的 1600 秒超时未生效")
		}
		if attempts < 3 {
			return errors.New("temporary")
		}
		return nil
	}})
	executor := &TaskExecutor{methods: registry}
	task := &taskbiz.TimedTask{MethodName: "migrated-retry", Params: []byte(`{"_execution":{"timeoutSeconds":1600,"maxRetryTimes":3,"retryIntervalSeconds":0}}`)}
	if err := executor.runMethod(context.Background(), task); err != nil || attempts != 3 {
		t.Fatalf("迁移重试未生效: attempts=%d err=%v", attempts, err)
	}
}

func TestRunMethodCancellationStopsRetryDelay(t *testing.T) {
	registry := platformtask.NewRegistry()
	started := make(chan struct{})
	registry.Register(platformtask.Method{Name: "cancel-retry", Run: func(context.Context, json.RawMessage) error { close(started); return errors.New("temporary") }})
	executor := &TaskExecutor{methods: registry}
	ctx, cancel := context.WithCancel(context.Background())
	done := make(chan error, 1)
	go func() {
		done <- executor.runMethod(ctx, &taskbiz.TimedTask{MethodName: "cancel-retry", Params: []byte(`{"_execution":{"maxRetryTimes":3,"retryIntervalSeconds":60}}`)})
	}()
	<-started
	cancel()
	select {
	case err := <-done:
		if !errors.Is(err, context.Canceled) {
			t.Fatalf("取消重试错误: %v", err)
		}
	case <-time.After(time.Second):
		t.Fatal("重试等待没有响应取消")
	}
}

func TestRunMethodCancellationWinsConcurrentAttemptFailure(t *testing.T) {
	for i := 0; i < 32; i++ {
		ctx, cancel := context.WithCancel(context.Background())
		registry := platformtask.NewRegistry()
		registry.Register(platformtask.Method{Name: "cancel-at-return", Run: func(context.Context, json.RawMessage) error { cancel(); return errors.New("temporary") }})
		err := (&TaskExecutor{methods: registry}).runMethod(ctx, &taskbiz.TimedTask{MethodName: "cancel-at-return", Params: []byte(`{"_execution":{"maxRetryTimes":3,"retryIntervalSeconds":60}}`)})
		cancel()
		if !errors.Is(err, context.Canceled) {
			t.Fatalf("并发取消被执行错误覆盖: %v", err)
		}
	}
}

func TestRunMethodUsesParentContext(t *testing.T) {
	const methodName = "worker-test-parent-context"
	started := make(chan struct{})
	methods := platformtask.NewRegistry()
	methods.Register(platformtask.Method{Name: methodName, Description: "test", Run: func(ctx context.Context, _ json.RawMessage) error {
		close(started)
		<-ctx.Done()
		return ctx.Err()
	}})
	executor := &TaskExecutor{methods: methods}
	ctx, cancel := context.WithCancel(context.Background())
	done := make(chan error, 1)
	go func() {
		done <- executor.runMethod(ctx, &taskbiz.TimedTask{MethodName: methodName})
	}()
	select {
	case <-started:
	case <-time.After(time.Second):
		t.Fatal("registered method did not start")
	}
	cancel()
	select {
	case err := <-done:
		if !errors.Is(err, context.Canceled) {
			t.Fatalf("runMethod() error = %v, want context.Canceled", err)
		}
	case <-time.After(time.Second):
		t.Fatal("runMethod() did not stop after parent cancellation")
	}
}

func TestRunMethodHonorsDeadlineWhenMethodIgnoresContext(t *testing.T) {
	const methodName = "worker-test-ignores-context"
	release := make(chan struct{})
	methods := platformtask.NewRegistry()
	methods.Register(platformtask.Method{Name: methodName, Description: "test", Run: func(context.Context, json.RawMessage) error {
		<-release
		return nil
	}})
	executor := &TaskExecutor{methods: methods}
	ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond)
	defer cancel()
	if err := executor.runMethod(ctx, &taskbiz.TimedTask{MethodName: methodName}); !errors.Is(err, errTaskTimeout) {
		close(release)
		t.Fatalf("runMethod() error = %v, want errTaskTimeout", err)
	}
	close(release)
}

func TestRunHandlesNilTask(t *testing.T) {
	executor := &TaskExecutor{}
	log := executor.Run(context.Background(), nil, "manual")
	if log == nil {
		t.Fatal("Run() returned nil log")
	}
	if log.Status != "fail" || log.ErrorMsg != "任务不能为空" {
		t.Fatalf("Run() log = %#v, want a failed nil-task result", log)
	}
}