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 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 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) } }