From fcf2ca869b20601902def8b1c58417975311d637 Mon Sep 17 00:00:00 2001 From: blackflytech Date: Wed, 1 Jul 2026 18:37:11 +0800 Subject: [PATCH] refactor: replace context.WithCancel with t.Context (#16509) ### Summary The addition of the Context method to Go's testing.T provides significant improvements for writing concurrent tests. It allows better management of goroutines, ensuring they properly exit and preventing issues like deadlocks and unfinished processes. By using Context, errors and cancellations can be handled more effectively, making tests more robust and easier to reason about. This change also enables tighter integration between tests and the application code, especially for systems that span multiple concurrent components. Overall, it simplifies test code and enhances test stability and maintainability. More info: [golang/go#18368](https://github.com/golang/go/issues/18368) Signed-off-by: blackflytech --- internal/agent/canvas/cancel_test.go | 3 +-- internal/harness/core/cancel_full_test.go | 21 +++++++-------------- internal/harness/core/react_graph_test.go | 3 +-- 3 files changed, 9 insertions(+), 18 deletions(-) diff --git a/internal/agent/canvas/cancel_test.go b/internal/agent/canvas/cancel_test.go index 5298425358..13eedc3304 100644 --- a/internal/agent/canvas/cancel_test.go +++ b/internal/agent/canvas/cancel_test.go @@ -47,8 +47,7 @@ func withCancelClient(t *testing.T) *miniredis.Miniredis { func TestWatchCancel_FiresAfterRequest(t *testing.T) { withCancelClient(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := t.Context() taskID := "task_test_1" fired := atomic.Bool{} diff --git a/internal/harness/core/cancel_full_test.go b/internal/harness/core/cancel_full_test.go index 303157dca3..519680741c 100644 --- a/internal/harness/core/cancel_full_test.go +++ b/internal/harness/core/cancel_full_test.go @@ -280,8 +280,7 @@ func TestBuildCancelFunc_CASFailStateDone(t *testing.T) { func TestDeriveAgentToolCancelContext(t *testing.T) { t.Run("Shallow/DoesNotPropagateSafePoint", func(t *testing.T) { parent := newCancelContext() - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := t.Context() child := parent.deriveAgentToolCancelContext(ctx) defer child.markDone() parent.triggerCancel(CancelAfterChatModel) @@ -293,8 +292,7 @@ func TestDeriveAgentToolCancelContext(t *testing.T) { }) t.Run("Shallow/ImmediateDoesNotPropagate", func(t *testing.T) { parent := newCancelContext() - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := t.Context() child := parent.deriveAgentToolCancelContext(ctx) defer child.markDone() parent.triggerImmediate() @@ -306,8 +304,7 @@ func TestDeriveAgentToolCancelContext(t *testing.T) { }) t.Run("Shallow/GrandchildNoPropagation", func(t *testing.T) { a := newCancelContext() - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := t.Context() b := a.deriveAgentToolCancelContext(ctx) c := b.deriveAgentToolCancelContext(ctx) t.Cleanup(func() { c.markDone(); b.markDone() }) @@ -370,8 +367,7 @@ func TestDeriveAgentToolCancelContext(t *testing.T) { }) t.Run("Recursive/GrandchildPropagation", func(t *testing.T) { a := newCancelContext() - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := t.Context() b := a.deriveAgentToolCancelContext(ctx) c := b.deriveAgentToolCancelContext(ctx) t.Cleanup(func() { c.markDone(); b.markDone() }) @@ -428,8 +424,7 @@ func TestDeriveAgentToolCancelContext_Race(t *testing.T) { }) t.Run("ChildCompletesBeforeEscalation", func(t *testing.T) { parent := newCancelContext() - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := t.Context() child := parent.deriveAgentToolCancelContext(ctx) parent.triggerCancel(CancelAfterChatModel) time.Sleep(50 * time.Millisecond) @@ -444,8 +439,7 @@ func TestDeriveAgentToolCancelContext_Race(t *testing.T) { }) t.Run("MultipleChildren_PartialCompletion", func(t *testing.T) { parent := newCancelContext() - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := t.Context() child1 := parent.deriveAgentToolCancelContext(ctx) child2 := parent.deriveAgentToolCancelContext(ctx) parent.triggerCancel(CancelAfterChatModel) @@ -1271,8 +1265,7 @@ func TestAgentCancelFunc_MultiCall_TimeoutDeadlineJoinUsesAbsoluteTime(t *testin func TestCancelContext_RecursiveGraceBoundary(t *testing.T) { a := newCancelContext() - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := t.Context() b := a.deriveAgentToolCancelContext(ctx) c := b.deriveAgentToolCancelContext(ctx) t.Cleanup(func() { c.markDone(); b.markDone() }) diff --git a/internal/harness/core/react_graph_test.go b/internal/harness/core/react_graph_test.go index 451b87b070..b0fcceabdd 100644 --- a/internal/harness/core/react_graph_test.go +++ b/internal/harness/core/react_graph_test.go @@ -82,8 +82,7 @@ func TestReActGraph_StreamWithInterrupt(t *testing.T) { t.Fatalf("NewReActGraph: %v", err) } - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + ctx := t.Context() outputCh, errCh := rg.Stream(ctx, &AgentInput{ Messages: []*schema.Message{schema.UserMessage("test")}}, nil, types.StreamModeValues)