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 <blackflytech@outlook.com>
This commit is contained in:
blackflytech
2026-07-01 18:37:11 +08:00
committed by GitHub
parent 5ba25a5267
commit fcf2ca869b
3 changed files with 9 additions and 18 deletions

View File

@@ -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{}

View File

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

View File

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