优化

internal/worker/task_executor_test.go

@@ -1,8 +1,15 @@
 package worker
 
 import (
+	"context"
+	"encoding/json"
+	"errors"
 	"net"
+	"strings"
 	"testing"
+	"time"
+
+	"kra/internal/biz"
 )
 
 func TestPrivateIP(t *testing.T) {
@@ -15,3 +22,76 @@ func TestPrivateIP(t *testing.T) {
 		t.Fatal("public IP was classified as private")
 	}
 }
+
+func TestRunMethodReportsParentDeadlineAsTimeout(t *testing.T) {
+	const methodName = "worker-test-parent-deadline"
+	biz.RegisterTaskMethod(methodName, "test", func(ctx context.Context, _ json.RawMessage) error {
+		<-ctx.Done()
+		return ctx.Err()
+	})
+	executor := &TaskExecutor{}
+	ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond)
+	defer cancel()
+	if err := executor.runMethod(ctx, &biz.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"
+	biz.RegisterTaskMethod(methodName, "test", func(context.Context, json.RawMessage) error {
+		panic("boom")
+	})
+	executor := &TaskExecutor{}
+	err := executor.runMethod(context.Background(), &biz.TimedTask{MethodName: methodName})
+	if err == nil || !strings.Contains(err.Error(), "panic: boom") {
+		t.Fatalf("runMethod() error = %v, want recovered panic", err)
+	}
+}
+
+func TestRunMethodUsesParentContext(t *testing.T) {
+	const methodName = "worker-test-parent-context"
+	started := make(chan struct{})
+	biz.RegisterTaskMethod(methodName, "test", func(ctx context.Context, _ json.RawMessage) error {
+		close(started)
+		<-ctx.Done()
+		return ctx.Err()
+	})
+	executor := &TaskExecutor{}
+	ctx, cancel := context.WithCancel(context.Background())
+	done := make(chan error, 1)
+	go func() {
+		done <- executor.runMethod(ctx, &biz.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{})
+	biz.RegisterTaskMethod(methodName, "test", func(context.Context, json.RawMessage) error {
+		<-release
+		return nil
+	})
+	executor := &TaskExecutor{}
+	ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond)
+	defer cancel()
+	if err := executor.runMethod(ctx, &biz.TimedTask{MethodName: methodName}); !errors.Is(err, errTaskTimeout) {
+		close(release)
+		t.Fatalf("runMethod() error = %v, want errTaskTimeout", err)
+	}
+	close(release)
+}