diff --git a/internal/agent/component/loop_test.go b/internal/agent/component/loop_test.go index 2154493661..119e55998b 100644 --- a/internal/agent/component/loop_test.go +++ b/internal/agent/component/loop_test.go @@ -27,7 +27,6 @@ package component import ( - "context" "testing" ) @@ -84,7 +83,8 @@ func TestLoop_InvokeIsNoOp(t *testing.T) { {"variable": "counter", "input_mode": "constant", "value": 7, "type": "number"}, }, }) - out, err := c.Invoke(context.Background(), nil, map[string]any{"in": 1}) + ctx := t.Context() + out, err := c.Invoke(ctx, nil, map[string]any{"in": 1}) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -97,7 +97,8 @@ func TestLoop_InvokeIsNoOp(t *testing.T) { // empty-map chunk and closes. func TestLoop_StreamMirrorsInvoke(t *testing.T) { c := NewLoopComponent(loopParam{}) - ch, err := c.Stream(context.Background(), nil, nil) + ctx := t.Context() + ch, err := c.Stream(ctx, nil, nil) if err != nil { t.Fatalf("Stream: %v", err) } diff --git a/internal/agent/component/memory_save_test.go b/internal/agent/component/memory_save_test.go index a51ad440f6..837a048e15 100644 --- a/internal/agent/component/memory_save_test.go +++ b/internal/agent/component/memory_save_test.go @@ -28,7 +28,8 @@ import ( func TestStubMemorySaver_DefaultReturnsError(t *testing.T) { SetMemorySaver(nil) saver := GetMemorySaver() - err := saver.Save(context.Background(), MemorySaveRequest{ + ctx := t.Context() + err := saver.Save(ctx, MemorySaveRequest{ MemoryIDs: []string{"m1"}, AgentID: "a1", }) @@ -48,7 +49,8 @@ func TestSetMemorySaver_Roundtrip(t *testing.T) { if got != custom { t.Fatalf("saver not registered") } - if err := got.Save(context.Background(), MemorySaveRequest{ + ctx := t.Context() + if err := got.Save(ctx, MemorySaveRequest{ MemoryIDs: []string{"m1"}, AgentResponse: "hi", }); err != nil { diff --git a/internal/agent/component/message_phase8b_test.go b/internal/agent/component/message_phase8b_test.go index 84fa69f06a..bdf5a055b2 100644 --- a/internal/agent/component/message_phase8b_test.go +++ b/internal/agent/component/message_phase8b_test.go @@ -49,7 +49,7 @@ func TestMessage_OutputFormatParam(t *testing.T) { "output_format": "html", }) state := canvas.NewCanvasState("r1", "t1") - ctx := withStateForTest(context.Background(), state) + ctx := withStateForTest(t.Context(), state) out, err := c.Invoke(ctx, nil, map[string]any{"text": "hello", "stream": false}) if err != nil { @@ -69,7 +69,7 @@ func TestMessage_OutputFormatInputOverride(t *testing.T) { "output_format": "html", }) state := canvas.NewCanvasState("r1", "t1") - ctx := withStateForTest(context.Background(), state) + ctx := withStateForTest(t.Context(), state) out, err := c.Invoke(ctx, nil, map[string]any{ "text": "hi", @@ -89,7 +89,7 @@ func TestMessage_OutputFormatInputOverride(t *testing.T) { func TestMessage_DownloadsExtraction(t *testing.T) { c, _ := NewMessageComponent(map[string]any{"text": "see attachment"}) state := canvas.NewCanvasState("r1", "t1") - ctx := withStateForTest(context.Background(), state) + ctx := withStateForTest(t.Context(), state) dl := map[string]any{ "doc_id": "d-1", @@ -124,7 +124,7 @@ func TestMessage_DownloadsExtraction(t *testing.T) { func TestMessage_DownloadJSONStringSuppressesContent(t *testing.T) { c, _ := NewMessageComponent(map[string]any{"text": "unused"}) state := canvas.NewCanvasState("r1", "t1") - ctx := withStateForTest(context.Background(), state) + ctx := withStateForTest(t.Context(), state) downloadJSON := `{"doc_id":"d-1","filename":"report.md","mime_type":"text/markdown","url":"/api/v1/agents/attachments/d-1/download","include_download_info_in_content":true}` out, err := c.Invoke(ctx, nil, map[string]any{ @@ -159,7 +159,7 @@ func TestMessage_AutoPlay_NoEngine(t *testing.T) { "auto_play": true, }) state := canvas.NewCanvasState("r1", "t1") - ctx := withStateForTest(context.Background(), state) + ctx := withStateForTest(t.Context(), state) out, err := c.Invoke(ctx, nil, map[string]any{"text": "hi", "stream": false}) if err != nil { @@ -192,7 +192,7 @@ func TestMessage_AutoPlay_Success(t *testing.T) { "lang": "en", }) state := canvas.NewCanvasState("r1", "t1") - ctx := withStateForTest(context.Background(), state) + ctx := withStateForTest(t.Context(), state) out, err := c.Invoke(ctx, nil, map[string]any{"text": "hi", "stream": false}) if err != nil { @@ -224,7 +224,7 @@ func TestMessage_MemorySave_NoService(t *testing.T) { c, _ := NewMessageComponent(map[string]any{"text": "hi"}) state := canvas.NewCanvasState("run-x", "task-x") state.Sys["query"] = "what?" - ctx := withStateForTest(context.Background(), state) + ctx := withStateForTest(t.Context(), state) out, err := c.Invoke(ctx, nil, map[string]any{ "text": "hi", @@ -260,7 +260,7 @@ func TestMessage_MemorySave_Success(t *testing.T) { state.Sys["canvas_id"] = "canvas-y" state.Sys["session_id"] = "session-y" state.Sys["agent_id"] = "agent-y" - ctx := withStateForTest(context.Background(), state) + ctx := withStateForTest(t.Context(), state) _, err := c.Invoke(ctx, nil, map[string]any{ "text": "hi", @@ -300,7 +300,7 @@ func TestMessage_MemorySave_FallbackIDs(t *testing.T) { c, _ := NewMessageComponent(map[string]any{"text": "hi"}) state := canvas.NewCanvasState("run-fallback", "task-fallback") state.Sys["query"] = "what?" - ctx := withStateForTest(context.Background(), state) + ctx := withStateForTest(t.Context(), state) _, err := c.Invoke(ctx, nil, map[string]any{ "text": "hi", @@ -338,7 +338,7 @@ func TestMessage_MemorySave_FromDSLParams(t *testing.T) { }) state := canvas.NewCanvasState("run-dsl", "task-dsl") state.Sys["query"] = "hello?" - ctx := withStateForTest(context.Background(), state) + ctx := withStateForTest(t.Context(), state) // Inputs simulate what the pipeline actually provides: only upstream // outputs, NO memory_ids or memory_save keys. @@ -377,7 +377,7 @@ func TestMessage_MemorySave_UserIDVariable(t *testing.T) { state := canvas.NewCanvasState("run-uid", "task-uid") state.Sys["query"] = "hello?" state.SetVar("begin", "user_id", "resolved-user-123") - ctx := withStateForTest(context.Background(), state) + ctx := withStateForTest(t.Context(), state) _, err := c.Invoke(ctx, nil, map[string]any{ "text": "hi", @@ -410,7 +410,7 @@ func TestMessage_MemorySave_UserIDLiteral(t *testing.T) { }) state := canvas.NewCanvasState("run-uid2", "task-uid2") state.Sys["query"] = "hello?" - ctx := withStateForTest(context.Background(), state) + ctx := withStateForTest(t.Context(), state) _, err := c.Invoke(ctx, nil, map[string]any{ "text": "hi", diff --git a/internal/agent/component/message_test.go b/internal/agent/component/message_test.go index d22bb744ee..ca86c7227c 100644 --- a/internal/agent/component/message_test.go +++ b/internal/agent/component/message_test.go @@ -32,7 +32,7 @@ func TestMessage_ResolveTemplate(t *testing.T) { c, _ := NewMessageComponent(nil) state := canvas.NewCanvasState("run-1", "task-1") state.Sys["query"] = "world" - ctx := withStateForTest(context.Background(), state) + ctx := withStateForTest(t.Context(), state) out, err := c.Invoke(ctx, nil, map[string]any{ "text": "hello {{sys.query}}", @@ -57,7 +57,7 @@ func TestMessage_ResolveListReferenceAsJSON(t *testing.T) { c, _ := NewMessageComponent(nil) state := canvas.NewCanvasState("run-list", "task-list") state.SetVar("list_0", "result", []any{"user: 1"}) - ctx := withStateForTest(context.Background(), state) + ctx := withStateForTest(t.Context(), state) out, err := c.Invoke(ctx, nil, map[string]any{ "text": "{{list_0@result}}", @@ -78,7 +78,7 @@ func TestMessage_Stream(t *testing.T) { c, _ := NewMessageComponent(nil) state := canvas.NewCanvasState("run-2", "task-2") state.Sys["query"] = "alice" - ctx := withStateForTest(context.Background(), state) + ctx := withStateForTest(t.Context(), state) ch, err := c.Stream(ctx, nil, map[string]any{ "text": "hi", @@ -107,7 +107,7 @@ func TestMessage_Stream(t *testing.T) { func TestMessage_NoTemplate(t *testing.T) { c, _ := NewMessageComponent(nil) state := canvas.NewCanvasState("run-3", "task-3") - ctx := withStateForTest(context.Background(), state) + ctx := withStateForTest(t.Context(), state) out, err := c.Invoke(ctx, nil, map[string]any{"text": "no refs here", "stream": false}) if err != nil { @@ -121,7 +121,7 @@ func TestMessage_NoTemplate(t *testing.T) { func TestMessage_RuntimeContentInput(t *testing.T) { c, _ := NewMessageComponent(nil) state := canvas.NewCanvasState("run-4", "task-4") - ctx := withStateForTest(context.Background(), state) + ctx := withStateForTest(t.Context(), state) out, err := c.Invoke(ctx, nil, map[string]any{"content": "from upstream", "stream": false}) if err != nil { @@ -135,7 +135,7 @@ func TestMessage_RuntimeContentInput(t *testing.T) { func TestMessage_EmitsAgentMessage(t *testing.T) { c, _ := NewMessageComponent(nil) state := canvas.NewCanvasState("run-emit", "task-emit") - ctx := withStateForTest(context.Background(), state) + ctx := withStateForTest(t.Context(), state) var emitted []string ctx = runtime.WithAgentMessageEmitter(ctx, func(contentDelta, thinkingDelta string) { if contentDelta != "" { @@ -164,7 +164,7 @@ func TestMessage_EmitsAgentMessage(t *testing.T) { func TestMessage_EmitsDirectCanvasMessage(t *testing.T) { c, _ := NewMessageComponent(nil) state := canvas.NewCanvasState("run-direct", "task-direct") - ctx := withStateForTest(context.Background(), state) + ctx := withStateForTest(t.Context(), state) var direct []string var agent []string ctx = runtime.WithAgentMessageEmitter(ctx, func(contentDelta, thinkingDelta string) { @@ -198,7 +198,7 @@ func TestMessage_NormalTemplateEmitsOnlyRenderedMessage(t *testing.T) { c, _ := NewMessageComponent(nil) state := canvas.NewCanvasState("run-normal-template", "task-normal-template") state.Sys["query"] = "world" - ctx := withStateForTest(context.Background(), state) + ctx := withStateForTest(t.Context(), state) var direct []string ctx = runtime.WithCanvasMessageEmitter(ctx, func(content string) { direct = append(direct, content) @@ -224,7 +224,7 @@ func TestMessage_NormalTemplateEmitsOnlyRenderedMessage(t *testing.T) { func TestMessage_SkipsEmissionWhenAgentAlreadyStreamed(t *testing.T) { c, _ := NewMessageComponent(nil) state := canvas.NewCanvasState("run-skip", "task-skip") - ctx := withStateForTest(context.Background(), state) + ctx := withStateForTest(t.Context(), state) var emitted []string ctx = runtime.WithAgentMessageEmitter(ctx, func(contentDelta, thinkingDelta string) { if contentDelta != "" { @@ -261,7 +261,7 @@ func TestMessage_SkipsEmissionWhenAgentAlreadyStreamed(t *testing.T) { func TestMessage_EmitsContentDifferentFromAgentStream(t *testing.T) { c, _ := NewMessageComponent(nil) state := canvas.NewCanvasState("run-distinct", "task-distinct") - ctx := withStateForTest(context.Background(), state) + ctx := withStateForTest(t.Context(), state) var emitted []string ctx = runtime.WithAgentMessageEmitter(ctx, func(contentDelta, thinkingDelta string) { if contentDelta != "" { @@ -301,7 +301,7 @@ func TestMessage_ConsumesDeferredAgentStream(t *testing.T) { return map[string]any{"content": "hello world"}, nil }, }) - ctx := withStateForTest(context.Background(), state) + ctx := withStateForTest(t.Context(), state) var emitted []string ctx = runtime.WithCanvasMessageEmitter(ctx, func(content string) { if content != "" { @@ -331,7 +331,7 @@ func TestMessage_DeferredStreamThinkingEvents(t *testing.T) { return map[string]any{"content": "answer"}, nil }, }) - ctx := withStateForTest(context.Background(), state) + ctx := withStateForTest(t.Context(), state) var events []string ctx = runtime.WithCanvasMessageEventEmitter(ctx, func(content string, startToThink, endToThink bool) { switch { @@ -361,7 +361,7 @@ func TestMessage_DeferredStreamUsesCompletedContent(t *testing.T) { return map[string]any{"content": "grounded answer [ID:1]"}, nil }, }) - ctx := withStateForTest(context.Background(), state) + ctx := withStateForTest(t.Context(), state) var streamed []string ctx = runtime.WithCanvasMessageEmitter(ctx, func(content string) { streamed = append(streamed, content) @@ -385,7 +385,7 @@ func TestMessage_DeferredStreamUsesCompletedContent(t *testing.T) { func TestMessage_FormalizedContentFallback(t *testing.T) { c, _ := NewMessageComponent(nil) state := canvas.NewCanvasState("run-5", "task-5") - ctx := withStateForTest(context.Background(), state) + ctx := withStateForTest(t.Context(), state) out, err := c.Invoke(ctx, nil, map[string]any{ "formalized_content": "retrieved answer", @@ -403,7 +403,7 @@ func TestMessage_FormalizedContentFallback(t *testing.T) { func TestMessage_SingleStringFallback(t *testing.T) { c, _ := NewMessageComponent(nil) state := canvas.NewCanvasState("run-6", "task-6") - ctx := withStateForTest(context.Background(), state) + ctx := withStateForTest(t.Context(), state) out, err := c.Invoke(ctx, nil, map[string]any{ "value": "single upstream text", diff --git a/internal/agent/component/parallel_test.go b/internal/agent/component/parallel_test.go index 507966854a..588e1c5bc4 100644 --- a/internal/agent/component/parallel_test.go +++ b/internal/agent/component/parallel_test.go @@ -34,7 +34,6 @@ package component import ( - "context" "slices" "testing" ) @@ -75,7 +74,9 @@ func TestParallel_InvokeIsNoOp(t *testing.T) { ItemsRef: "sys.arr", MaxConcurrency: 3, }) - out, err := c.Invoke(context.Background(), nil, map[string]any{"in": 1}) + ctx := t.Context() + + out, err := c.Invoke(ctx, nil, map[string]any{"in": 1}) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -88,7 +89,9 @@ func TestParallel_InvokeIsNoOp(t *testing.T) { // empty-map chunk and closes. func TestParallel_StreamMirrorsInvoke(t *testing.T) { c := NewParallelComponent(ParallelParam{}) - ch, err := c.Stream(context.Background(), nil, nil) + ctx := t.Context() + + ch, err := c.Stream(ctx, nil, nil) if err != nil { t.Fatalf("Stream: %v", err) } diff --git a/internal/agent/component/sampler_params_test.go b/internal/agent/component/sampler_params_test.go index 90284b6233..c407155480 100644 --- a/internal/agent/component/sampler_params_test.go +++ b/internal/agent/component/sampler_params_test.go @@ -35,7 +35,9 @@ func TestLLM_ForwardsTopP(t *testing.T) { ModelID: "echo", TopP: &topP, }) - if _, err := c.Invoke(context.Background(), nil, map[string]any{"user_prompt": "hi"}); err != nil { + ctx := t.Context() + + if _, err := c.Invoke(ctx, nil, map[string]any{"user_prompt": "hi"}); err != nil { t.Fatalf("Invoke: %v", err) } if stub.calls != 1 { @@ -57,9 +59,11 @@ func TestLLM_ForwardsTopP(t *testing.T) { func TestLLM_TopPFromInputs(t *testing.T) { stub := &stubInvoker{resp: &ChatInvokeResponse{Content: "ok", Model: "echo"}} withStubInvoker(t, stub) + ctx := t.Context() c := NewLLMComponent(LLMParam{ModelID: "echo"}) - if _, err := c.Invoke(context.Background(), nil, map[string]any{ + + if _, err := c.Invoke(ctx, nil, map[string]any{ "user_prompt": "hi", "top_p": 0.7, }); err != nil { @@ -78,9 +82,10 @@ func TestLLM_TopPFromInputs(t *testing.T) { func TestLLM_NoTopPByDefault(t *testing.T) { stub := &stubInvoker{resp: &ChatInvokeResponse{Content: "ok", Model: "echo"}} withStubInvoker(t, stub) + ctx := t.Context() c := NewLLMComponent(LLMParam{ModelID: "echo"}) - if _, err := c.Invoke(context.Background(), nil, map[string]any{"user_prompt": "hi"}); err != nil { + if _, err := c.Invoke(ctx, nil, map[string]any{"user_prompt": "hi"}); err != nil { t.Fatalf("Invoke: %v", err) } if stub.captured == nil { @@ -128,6 +133,7 @@ func TestAgentParam_ForwardsTopP(t *testing.T) { } return &schema.Message{Content: "ok"}, nil }) + ctx := t.Context() topP := 0.5 c := NewAgentComponent(AgentParam{ @@ -135,7 +141,7 @@ func TestAgentParam_ForwardsTopP(t *testing.T) { TopP: &topP, MaxRounds: 1, }) - if _, err := c.Invoke(context.Background(), nil, map[string]any{"user_prompt": "hi"}); err != nil { + if _, err := c.Invoke(ctx, nil, map[string]any{"user_prompt": "hi"}); err != nil { t.Fatalf("Invoke: %v", err) } } @@ -148,9 +154,10 @@ func TestAgent_TopPFromInputs(t *testing.T) { } return &schema.Message{Content: "ok"}, nil }) + ctx := t.Context() c := NewAgentComponent(AgentParam{ModelID: "echo", MaxRounds: 1}) - if _, err := c.Invoke(context.Background(), nil, map[string]any{ + if _, err := c.Invoke(ctx, nil, map[string]any{ "user_prompt": "hi", "top_p": 0.42, }); err != nil { diff --git a/internal/agent/component/stagehand_runtime_integration_test.go b/internal/agent/component/stagehand_runtime_integration_test.go index 2062e560e5..433e574fcc 100644 --- a/internal/agent/component/stagehand_runtime_integration_test.go +++ b/internal/agent/component/stagehand_runtime_integration_test.go @@ -129,7 +129,7 @@ func TestStagehandRuntime_Extract(t *testing.T) { Schema: schema, } - ctx, cancel := context.WithTimeout(context.Background(), 3*time.Minute) + ctx, cancel := context.WithTimeout(t.Context(), 3*time.Minute) defer cancel() t.Logf("starting stagehand RunExtract (timeout 3m); spawns subprocess, calls LLM once with schema=%s", @@ -240,7 +240,7 @@ func TestBrowser_E2E_Extract(t *testing.T) { t.Fatalf("NewBrowserComponent: %v", err) } - ctx := canvas.WithState(context.Background(), canvas.NewCanvasState("run-1", "task-1")) + ctx := canvas.WithState(t.Context(), canvas.NewCanvasState("run-1", "task-1")) state, _, _ := runtime.GetStateFromContext[*runtime.CanvasState](ctx) state.Sys["user_id"] = "tenant-1" diff --git a/internal/agent/component/stagehand_runtime_test.go b/internal/agent/component/stagehand_runtime_test.go index b6943f25a0..056455c485 100644 --- a/internal/agent/component/stagehand_runtime_test.go +++ b/internal/agent/component/stagehand_runtime_test.go @@ -46,6 +46,7 @@ func cacheSize(r *stagehandRuntime) int { func TestStagehandRuntime_ValidatesRequiredFields(t *testing.T) { r := newStagehandRuntime(time.Hour, 0, time.Minute) // TTL large → no sweeper interference t.Cleanup(func() { _ = r.Close() }) + ctx := t.Context() cases := []struct { name string @@ -58,7 +59,7 @@ func TestStagehandRuntime_ValidatesRequiredFields(t *testing.T) { } for _, tc := range cases { t.Run(tc.name, func(t *testing.T) { - _, err := r.RunTask(context.Background(), tc.req) + _, err := r.RunTask(ctx, tc.req) if err == nil { t.Fatalf("expected error for %s, got nil", tc.name) } @@ -535,7 +536,8 @@ func TestStagehandRuntime_SetDefaultStagehandInvoker(t *testing.T) { if got == nil { t.Fatal("getDefaultStagehandInvoker returned nil after swap") } - out, err := got.RunTask(context.Background(), RunTaskRequest{Instruction: "x", ModelName: "m", APIKey: "k"}) + ctx := t.Context() + out, err := got.RunTask(ctx, RunTaskRequest{Instruction: "x", ModelName: "m", APIKey: "k"}) if err != nil { t.Fatalf("RunTask: %v", err) } diff --git a/internal/agent/component/streaming_test.go b/internal/agent/component/streaming_test.go index 3128c673d9..e76c00ac9e 100644 --- a/internal/agent/component/streaming_test.go +++ b/internal/agent/component/streaming_test.go @@ -27,9 +27,10 @@ import ( func TestLLM_Stream_HappyPath(t *testing.T) { stub := &stubInvoker{resp: &ChatInvokeResponse{Content: "hello", Model: "echo"}} withStubInvoker(t, stub) + ctx := t.Context() c := NewLLMComponent(LLMParam{ModelID: "echo"}) - ch, err := c.Stream(context.Background(), nil, map[string]any{"user_prompt": "hi"}) + ch, err := c.Stream(ctx, nil, map[string]any{"user_prompt": "hi"}) if err != nil { t.Fatalf("Stream: %v", err) } @@ -61,9 +62,10 @@ func TestLLM_Stream_HappyPath(t *testing.T) { func TestLLM_Stream_Error(t *testing.T) { stub := &stubInvoker{err: context.DeadlineExceeded} withStubInvoker(t, stub) + ctx := t.Context() c := NewLLMComponent(LLMParam{ModelID: "echo"}) - ch, err := c.Stream(context.Background(), nil, map[string]any{"user_prompt": "hi"}) + ch, err := c.Stream(ctx, nil, map[string]any{"user_prompt": "hi"}) if err != nil { t.Fatalf("Stream: %v", err) } @@ -86,7 +88,7 @@ func TestLLM_Stream_RespectsCancellation(t *testing.T) { withStubInvoker(t, stub) c := NewLLMComponent(LLMParam{ModelID: "echo"}) - ctx, cancel := context.WithCancel(context.Background()) + ctx, cancel := context.WithCancel(t.Context()) cancel() // pre-cancel ch, err := c.Stream(ctx, nil, map[string]any{"user_prompt": "hi"}) @@ -111,9 +113,10 @@ func TestLLM_Stream_RespectsCancellation(t *testing.T) { func TestLLM_Stream_BufferDoesNotBlock(t *testing.T) { stub := &stubInvoker{resp: &ChatInvokeResponse{Content: "ok", Model: "echo"}} withStubInvoker(t, stub) + ctx := t.Context() c := NewLLMComponent(LLMParam{ModelID: "echo"}) - ch, err := c.Stream(context.Background(), nil, map[string]any{"user_prompt": "hi"}) + ch, err := c.Stream(ctx, nil, map[string]any{"user_prompt": "hi"}) if err != nil { t.Fatalf("Stream: %v", err) } diff --git a/internal/agent/component/string_transform_test.go b/internal/agent/component/string_transform_test.go index 3f2cde17a8..2d1ac5f30d 100644 --- a/internal/agent/component/string_transform_test.go +++ b/internal/agent/component/string_transform_test.go @@ -17,7 +17,6 @@ package component import ( - "context" "reflect" "testing" @@ -34,7 +33,7 @@ func TestStringTransform_SplitBasic(t *testing.T) { t.Fatalf("NewStringTransformComponent: %v", err) } state := canvas.NewCanvasState("run-1", "task-1") - ctx := canvas.WithState(context.Background(), state) + ctx := canvas.WithState(t.Context(), state) out, err := c.Invoke(ctx, nil, map[string]any{"line": "a,b;c"}) if err != nil { @@ -54,7 +53,7 @@ func TestStringTransform_SplitNoDelim(t *testing.T) { "delimiters": []string{","}, }) state := canvas.NewCanvasState("run-2", "task-2") - ctx := canvas.WithState(context.Background(), state) + ctx := canvas.WithState(t.Context(), state) out, err := c.Invoke(ctx, nil, map[string]any{"line": "abc"}) if err != nil { @@ -75,7 +74,7 @@ func TestStringTransform_Merge(t *testing.T) { "script": "{{x}} and {{y}}", }) state := canvas.NewCanvasState("run-3", "task-3") - ctx := canvas.WithState(context.Background(), state) + ctx := canvas.WithState(t.Context(), state) out, err := c.Invoke(ctx, nil, map[string]any{"x": "foo", "y": "bar"}) if err != nil { @@ -95,7 +94,7 @@ func TestStringTransform_MergeIterationAliases(t *testing.T) { state := canvas.NewCanvasState("run-iter", "task-iter") state.Globals["__item__"] = "beta" state.Globals["__index__"] = 1 - ctx := canvas.WithState(context.Background(), state) + ctx := canvas.WithState(t.Context(), state) out, err := c.Invoke(ctx, nil, map[string]any{}) if err != nil { @@ -116,7 +115,7 @@ func TestStringTransform_SplitFromStateRef(t *testing.T) { }) state := canvas.NewCanvasState("run-4", "task-4") state.Outputs["cpn_0"] = map[string]any{"x": "alpha,beta,gamma"} - ctx := canvas.WithState(context.Background(), state) + ctx := canvas.WithState(t.Context(), state) out, err := c.Invoke(ctx, nil, nil) if err != nil { @@ -138,7 +137,7 @@ func TestStringTransform_MergeMissingPlaceholder(t *testing.T) { "script": "hello {{name}}", }) state := canvas.NewCanvasState("run-5", "task-5") - ctx := canvas.WithState(context.Background(), state) + ctx := canvas.WithState(t.Context(), state) out, err := c.Invoke(ctx, nil, map[string]any{}) if err != nil { diff --git a/internal/agent/component/structured_output_test.go b/internal/agent/component/structured_output_test.go index 1e30575158..bc7002a7f3 100644 --- a/internal/agent/component/structured_output_test.go +++ b/internal/agent/component/structured_output_test.go @@ -102,12 +102,13 @@ func TestLLM_Invoke_OutputStructure_ValidFirstTry(t *testing.T) { Model: "echo", }} withStubInvoker(t, stub) + ctx := t.Context() c := NewLLMComponent(LLMParam{ ModelID: "echo", OutputStructure: map[string]any{"name": "", "age": 0}, }) - out, err := c.Invoke(context.Background(), nil, map[string]any{"user_prompt": "who?"}) + out, err := c.Invoke(ctx, nil, map[string]any{"user_prompt": "who?"}) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -136,12 +137,13 @@ func TestLLM_Invoke_OutputStructure_RetryOnInvalid(t *testing.T) { onCall: func() { calls++ }, } withStubInvoker(t, inv) + ctx := t.Context() c := NewLLMComponent(LLMParam{ ModelID: "echo", OutputStructure: map[string]any{"name": ""}, }) - out, err := c.Invoke(context.Background(), nil, map[string]any{"user_prompt": "who?"}) + out, err := c.Invoke(ctx, nil, map[string]any{"user_prompt": "who?"}) if err != nil { t.Fatalf("Invoke: %v", err) } @@ -173,12 +175,13 @@ func TestLLM_Invoke_OutputStructure_RetryStillFails(t *testing.T) { onCall: func() { calls++ }, } withStubInvoker(t, inv) + ctx := t.Context() c := NewLLMComponent(LLMParam{ ModelID: "echo", OutputStructure: map[string]any{"x": 0}, }) - out, err := c.Invoke(context.Background(), nil, map[string]any{"user_prompt": "go"}) + out, err := c.Invoke(ctx, nil, map[string]any{"user_prompt": "go"}) if err != nil { t.Fatalf("Invoke should not error on parse failure: %v", err) } diff --git a/internal/agent/component/switch_test.go b/internal/agent/component/switch_test.go index 773d5e3bc6..b9109295ef 100644 --- a/internal/agent/component/switch_test.go +++ b/internal/agent/component/switch_test.go @@ -17,7 +17,6 @@ package component import ( - "context" "testing" "ragflow/internal/agent/canvas" @@ -56,7 +55,7 @@ func TestSwitch_AndMatches(t *testing.T) { s, _ := NewSwitchComponent(nil) state := canvas.NewCanvasState("run-1", "task-1") state.Sys["x"] = "yes" - ctx := withStateForTest(context.Background(), state) + ctx := withStateForTest(t.Context(), state) inputs := map[string]any{ "conditions": []any{ @@ -87,7 +86,7 @@ func TestSwitch_OrMatches(t *testing.T) { state := canvas.NewCanvasState("run-2", "task-2") state.Sys["score"] = "85" state.Sys["flag"] = "no" - ctx := withStateForTest(context.Background(), state) + ctx := withStateForTest(t.Context(), state) inputs := map[string]any{ "conditions": []any{ @@ -129,7 +128,7 @@ func TestSwitch_DefaultFallback(t *testing.T) { s, _ := NewSwitchComponent(nil) state := canvas.NewCanvasState("run-3", "task-3") state.Sys["x"] = "no" - ctx := withStateForTest(context.Background(), state) + ctx := withStateForTest(t.Context(), state) inputs := map[string]any{ "conditions": []any{ @@ -156,7 +155,7 @@ func TestSwitch_LegacyEndCpnIDsFallback(t *testing.T) { s, _ := NewSwitchComponent(nil) state := canvas.NewCanvasState("run-end-cpn", "task-end-cpn") state.Sys["x"] = "no" - ctx := withStateForTest(context.Background(), state) + ctx := withStateForTest(t.Context(), state) inputs := map[string]any{ "conditions": []any{ @@ -186,7 +185,7 @@ func TestSwitch_ContainsAndEmpty(t *testing.T) { state := canvas.NewCanvasState("run-4", "task-4") state.Sys["body"] = "hello world" state.Sys["opt"] = "" - ctx := withStateForTest(context.Background(), state) + ctx := withStateForTest(t.Context(), state) inputs := map[string]any{ "conditions": []any{ @@ -229,7 +228,7 @@ func TestSwitch_LegacyConditionsAndArrayTo(t *testing.T) { }) state := canvas.NewCanvasState("run-legacy", "task-legacy") state.SetVar("UserFillUp:Menu", "demo", "loop") - ctx := withStateForTest(context.Background(), state) + ctx := withStateForTest(t.Context(), state) out, err := s.Invoke(ctx, nil, nil) if err != nil { @@ -253,7 +252,7 @@ func TestSwitch_NilUpstreamContainsEmptyNeedleMatches(t *testing.T) { s, _ := NewSwitchComponent(nil) state := canvas.NewCanvasState("run-nil-contains", "task-nil-contains") state.Sys["answer"] = nil - ctx := withStateForTest(context.Background(), state) + ctx := withStateForTest(t.Context(), state) inputs := map[string]any{ "conditions": []any{ @@ -284,7 +283,7 @@ func TestSwitch_NilUpstreamContainsNonEmptyDoesNotMatch(t *testing.T) { s, _ := NewSwitchComponent(nil) state := canvas.NewCanvasState("run-nil-needle", "task-nil-needle") state.Sys["answer"] = nil - ctx := withStateForTest(context.Background(), state) + ctx := withStateForTest(t.Context(), state) inputs := map[string]any{ "conditions": []any{ @@ -316,7 +315,7 @@ func TestSwitch_NilValueContainsDoesNotRaise(t *testing.T) { s, _ := NewSwitchComponent(nil) state := canvas.NewCanvasState("run-nil-value", "task-nil-value") state.Sys["answer"] = "foobar" - ctx := withStateForTest(context.Background(), state) + ctx := withStateForTest(t.Context(), state) inputs := map[string]any{ "conditions": []any{ @@ -350,7 +349,7 @@ func TestSwitch_NilUpstreamStartWithEndWithDoNotCrash(t *testing.T) { s, _ := NewSwitchComponent(nil) state := canvas.NewCanvasState("run-nil-start-end", "task-nil-start-end") state.Sys["answer"] = nil - ctx := withStateForTest(context.Background(), state) + ctx := withStateForTest(t.Context(), state) for _, tc := range []struct { name string @@ -393,7 +392,7 @@ func TestSwitch_MultiTargetTo(t *testing.T) { s, _ := NewSwitchComponent(nil) state := canvas.NewCanvasState("run-multi", "task-multi") state.SetVar("UserFillUp:Menu", "demo", "data_ops") - ctx := withStateForTest(context.Background(), state) + ctx := withStateForTest(t.Context(), state) inputs := map[string]any{ "conditions": []any{ @@ -430,7 +429,7 @@ func TestSwitch_MultiTargetTo(t *testing.T) { func TestSwitch_EmptyAndConditionFallsThrough(t *testing.T) { s, _ := NewSwitchComponent(nil) state := canvas.NewCanvasState("run-empty-and", "task-1") - ctx := withStateForTest(context.Background(), state) + ctx := withStateForTest(t.Context(), state) // Empty clauses: must not match. Should fall through to default. inputs := map[string]any{ @@ -462,7 +461,7 @@ func TestSwitch_EmptyAndConditionFallsThrough(t *testing.T) { func TestSwitch_LegacyEmptyItemsFallsThrough(t *testing.T) { s, _ := NewSwitchComponent(nil) state := canvas.NewCanvasState("run-legacy-empty", "task-1") - ctx := withStateForTest(context.Background(), state) + ctx := withStateForTest(t.Context(), state) inputs := map[string]any{ "conditions": []any{ @@ -496,7 +495,7 @@ func TestSwitch_SatisfiedAndConditionStillRoutes(t *testing.T) { s, _ := NewSwitchComponent(nil) state := canvas.NewCanvasState("run-and-ok", "task-1") state.Sys["greeting"] = "hello world" - ctx := withStateForTest(context.Background(), state) + ctx := withStateForTest(t.Context(), state) inputs := map[string]any{ "conditions": []any{ diff --git a/internal/agent/component/tool_call_memory_test.go b/internal/agent/component/tool_call_memory_test.go index 30783db5e1..581253c4ee 100644 --- a/internal/agent/component/tool_call_memory_test.go +++ b/internal/agent/component/tool_call_memory_test.go @@ -31,8 +31,9 @@ import ( func TestAddToolCallMemory_NoToolCalls(t *testing.T) { stub := &stubInvoker{resp: &ChatInvokeResponse{Content: "ok", Model: "echo"}} withStubInvoker(t, stub) + ctx := t.Context() - got, err := addToolCallMemory(context.Background(), nil, AgentParam{ModelID: "echo"}, &schema.Message{Content: "no tools"}) + got, err := addToolCallMemory(ctx, nil, AgentParam{ModelID: "echo"}, &schema.Message{Content: "no tools"}) if err != nil { t.Fatalf("err: %v", err) } @@ -70,7 +71,7 @@ func TestAddToolCallMemory_SummarizesAndAppendsToState(t *testing.T) { state := runtime.NewCanvasState("rid", "tid") c := NewAgentComponent(AgentParam{ModelID: "echo", MaxRounds: 1}) - ctx := runtime.WithState(context.Background(), state) + ctx := runtime.WithState(t.Context(), state) _, err := c.Invoke(ctx, nil, map[string]any{"user_prompt": "do it"}) if err != nil { t.Fatalf("Invoke: %v", err) @@ -112,7 +113,7 @@ func TestAddToolCallMemory_LLMFailure(t *testing.T) { state := runtime.NewCanvasState("rid", "tid") c := NewAgentComponent(AgentParam{ModelID: "echo", MaxRounds: 1}) - ctx := runtime.WithState(context.Background(), state) + ctx := runtime.WithState(t.Context(), state) _, err := c.Invoke(ctx, nil, map[string]any{"user_prompt": "do it"}) if err != nil { t.Fatalf("Invoke should not error when memory summary fails: %v", err) diff --git a/internal/agent/component/tool_dispatch_test.go b/internal/agent/component/tool_dispatch_test.go index 75f16f1253..761e97873e 100644 --- a/internal/agent/component/tool_dispatch_test.go +++ b/internal/agent/component/tool_dispatch_test.go @@ -64,12 +64,13 @@ func TestPhase3_6_ToolDSLLoading(t *testing.T) { return &schema.Message{Role: schema.Assistant, Content: "ok"}, nil }) + ctx := t.Context() c := NewAgentComponent(AgentParam{ ModelID: "stub", MaxRounds: 1, Tools: []string{"retrieval"}, // known tool }) - _, err := c.Invoke(context.Background(), nil, map[string]any{ + _, err := c.Invoke(ctx, nil, map[string]any{ "user_prompt": "test", }) if err != nil { diff --git a/internal/agent/component/userfillup_test.go b/internal/agent/component/userfillup_test.go index 66ee37d23b..933bb8b46a 100644 --- a/internal/agent/component/userfillup_test.go +++ b/internal/agent/component/userfillup_test.go @@ -17,7 +17,6 @@ package component import ( - "context" "testing" "ragflow/internal/agent/canvas" @@ -33,7 +32,7 @@ func TestUserFillUp_RendersTips(t *testing.T) { "tips": "Hello {{name}}", }) state := canvas.NewCanvasState("run-1", "task-1") - ctx := withStateForTest(context.Background(), state) + ctx := withStateForTest(t.Context(), state) out, err := c.Invoke(ctx, nil, map[string]any{ "inputs": map[string]any{ @@ -57,7 +56,7 @@ func TestUserFillUp_DisableTips(t *testing.T) { "tips": "Should not render", }) state := canvas.NewCanvasState("run-2", "task-2") - ctx := withStateForTest(context.Background(), state) + ctx := withStateForTest(t.Context(), state) out, err := c.Invoke(ctx, nil, map[string]any{ "inputs": map[string]any{ @@ -83,7 +82,7 @@ func TestUserFillUp_DisableTips(t *testing.T) { func TestUserFillUp_PassesThroughInputs(t *testing.T) { c, _ := New(componentNameUserFillUp, map[string]any{"enable_tips": false}) state := canvas.NewCanvasState("run-3", "task-3") - ctx := withStateForTest(context.Background(), state) + ctx := withStateForTest(t.Context(), state) out, err := c.Invoke(ctx, nil, map[string]any{ "inputs": map[string]any{ @@ -116,7 +115,7 @@ func TestUserFillUp_FileInputStub(t *testing.T) { "tips": "Upload {{cv}} please", }) state := canvas.NewCanvasState("run-4", "task-4") - ctx := withStateForTest(context.Background(), state) + ctx := withStateForTest(t.Context(), state) out, err := c.Invoke(ctx, nil, map[string]any{ "inputs": map[string]any{ diff --git a/internal/agent/component/variable_aggregator_test.go b/internal/agent/component/variable_aggregator_test.go index d7e90bcf62..48fda966ff 100644 --- a/internal/agent/component/variable_aggregator_test.go +++ b/internal/agent/component/variable_aggregator_test.go @@ -17,7 +17,6 @@ package component import ( - "context" "testing" "ragflow/internal/agent/canvas" @@ -31,7 +30,7 @@ func TestVariableAggregator_FirstNonEmpty(t *testing.T) { state.Outputs["cpn_1"] = map[string]any{"y": "second-a"} state.Outputs["cpn_2"] = map[string]any{"y": "second-b"} state.Outputs["cpn_3"] = map[string]any{"y": "second-c"} - ctx := canvas.WithState(context.Background(), state) + ctx := canvas.WithState(t.Context(), state) groups := []map[string]any{ { @@ -84,7 +83,7 @@ func TestVariableAggregator_SkipsEmptyString(t *testing.T) { state := canvas.NewCanvasState("run-2", "task-2") state.Outputs["cpn_0"] = map[string]any{"x": ""} state.Outputs["cpn_1"] = map[string]any{"y": "picked"} - ctx := canvas.WithState(context.Background(), state) + ctx := canvas.WithState(t.Context(), state) groups := []map[string]any{ { @@ -115,7 +114,7 @@ func TestVariableAggregator_MultipleGroups(t *testing.T) { state.Sys["a"] = "alpha" state.Sys["b"] = "" state.Env["c"] = "gamma" - ctx := canvas.WithState(context.Background(), state) + ctx := canvas.WithState(t.Context(), state) groups := []map[string]any{ { @@ -163,7 +162,7 @@ func TestVariableAggregator_MultipleGroups(t *testing.T) { func TestVariableAggregator_AllEmpty(t *testing.T) { state := canvas.NewCanvasState("run-4", "task-4") state.Outputs["cpn_0"] = map[string]any{} - ctx := canvas.WithState(context.Background(), state) + ctx := canvas.WithState(t.Context(), state) groups := []map[string]any{ { diff --git a/internal/agent/component/variable_assigner_test.go b/internal/agent/component/variable_assigner_test.go index 27137b8b9e..f490b81e74 100644 --- a/internal/agent/component/variable_assigner_test.go +++ b/internal/agent/component/variable_assigner_test.go @@ -17,7 +17,6 @@ package component import ( - "context" "reflect" "testing" @@ -28,7 +27,7 @@ import ( func TestVariableAssigner_Append(t *testing.T) { state := canvas.NewCanvasState("run-1", "task-1") state.Outputs["cpn_0"] = map[string]any{"xs": []any{1, 2}} - ctx := canvas.WithState(context.Background(), state) + ctx := canvas.WithState(t.Context(), state) vars := []map[string]any{ { @@ -62,7 +61,7 @@ func TestVariableAssigner_Overwrite(t *testing.T) { state := canvas.NewCanvasState("run-2", "task-2") state.Outputs["cpn_0"] = map[string]any{"x": "old"} state.Outputs["cpn_1"] = map[string]any{"y": "fresh"} - ctx := canvas.WithState(context.Background(), state) + ctx := canvas.WithState(t.Context(), state) vars := []map[string]any{ { @@ -88,7 +87,7 @@ func TestVariableAssigner_Overwrite(t *testing.T) { func TestVariableAssigner_DivideByZero(t *testing.T) { state := canvas.NewCanvasState("run-3", "task-3") state.Outputs["cpn_0"] = map[string]any{"n": 6.0} - ctx := canvas.WithState(context.Background(), state) + ctx := canvas.WithState(t.Context(), state) vars := []map[string]any{ { @@ -133,7 +132,7 @@ func TestVariableAssigner_Clear(t *testing.T) { "c": map[string]any{"k": "v"}, "d": 42, } - ctx := canvas.WithState(context.Background(), state) + ctx := canvas.WithState(t.Context(), state) vars := []map[string]any{ {"variable": "cpn_0@a", "operator": "clear", "parameter": "x"}, @@ -166,7 +165,7 @@ func TestVariableAssigner_Clear(t *testing.T) { func TestVariableAssigner_Arithmetic(t *testing.T) { state := canvas.NewCanvasState("run-5", "task-5") state.Outputs["cpn_0"] = map[string]any{"n": 10.0} - ctx := canvas.WithState(context.Background(), state) + ctx := canvas.WithState(t.Context(), state) vars := []map[string]any{ {"variable": "cpn_0@n", "operator": "+=", "parameter": 5}, @@ -191,7 +190,7 @@ func TestVariableAssigner_Arithmetic(t *testing.T) { func TestVariableAssigner_RemoveFirstLast(t *testing.T) { state := canvas.NewCanvasState("run-6", "task-6") state.Outputs["cpn_0"] = map[string]any{"xs": []any{"a", "b", "c", "d"}} - ctx := canvas.WithState(context.Background(), state) + ctx := canvas.WithState(t.Context(), state) vars := []map[string]any{ {"variable": "cpn_0@xs", "operator": "remove_first", "parameter": "x"}, @@ -219,7 +218,7 @@ func TestVariableAssigner_RemoveFirstLast(t *testing.T) { // TestVariableAssigner_SysTarget: variable="sys.x" → state.Sys is written. func TestVariableAssigner_SysTarget(t *testing.T) { state := canvas.NewCanvasState("run-7", "task-7") - ctx := canvas.WithState(context.Background(), state) + ctx := canvas.WithState(t.Context(), state) vars := []map[string]any{ {"variable": "sys.x", "operator": "set", "parameter": "hello"}, diff --git a/internal/agent/component/vision_test.go b/internal/agent/component/vision_test.go index 1ee411f04a..c0ba352084 100644 --- a/internal/agent/component/vision_test.go +++ b/internal/agent/component/vision_test.go @@ -17,7 +17,6 @@ package component import ( - "context" "reflect" "testing" @@ -285,10 +284,11 @@ func TestBuildMessagesWithImages_WithImages_UsesUserInputMultiContent(t *testing func TestLLM_Invoke_ForwardsImagesToInvoker(t *testing.T) { stub := &stubInvoker{resp: &ChatInvokeResponse{Content: "ok", Model: "echo"}} withStubInvoker(t, stub) + ctx := t.Context() uri := "data:image/png;base64,iVBORw0KGgo=" c := NewLLMComponent(LLMParam{ModelID: "echo"}) - _, err := c.Invoke(context.Background(), nil, map[string]any{ + _, err := c.Invoke(ctx, nil, map[string]any{ "user_prompt": "what is this?", "visual_files": []string{uri}, }) @@ -318,9 +318,10 @@ func TestLLM_Invoke_ForwardsImagesToInvoker(t *testing.T) { func TestLLM_Invoke_NoVisualFiles_BackwardCompat(t *testing.T) { stub := &stubInvoker{resp: &ChatInvokeResponse{Content: "ok", Model: "echo"}} withStubInvoker(t, stub) + ctx := t.Context() c := NewLLMComponent(LLMParam{ModelID: "echo"}) - _, err := c.Invoke(context.Background(), nil, map[string]any{ + _, err := c.Invoke(ctx, nil, map[string]any{ "user_prompt": "hi", }) if err != nil { @@ -347,10 +348,11 @@ func TestLLM_Invoke_NoVisualFiles_BackwardCompat(t *testing.T) { func TestLLM_Invoke_VisualFilesAsString(t *testing.T) { stub := &stubInvoker{resp: &ChatInvokeResponse{Content: "ok", Model: "echo"}} withStubInvoker(t, stub) + ctx := t.Context() uri := "data:image/jpeg;base64,/9j/4AAQ" c := NewLLMComponent(LLMParam{ModelID: "echo"}) - _, err := c.Invoke(context.Background(), nil, map[string]any{ + _, err := c.Invoke(ctx, nil, map[string]any{ "user_prompt": "describe", "visual_files": "see " + uri, }) diff --git a/internal/agent/sandbox/e2b_test.go b/internal/agent/sandbox/e2b_test.go index bd794928bf..040ee61ed6 100644 --- a/internal/agent/sandbox/e2b_test.go +++ b/internal/agent/sandbox/e2b_test.go @@ -79,7 +79,7 @@ func TestE2BProvider_Initialize_MissingCreds(t *testing.T) { t.Setenv(k, "") } p := newE2BProviderFromEnv() - err := p.Initialize(context.Background()) + err := p.Initialize(t.Context()) if err == nil { t.Fatalf("Initialize with no creds: got nil error, want one") } @@ -98,7 +98,7 @@ func TestE2BProvider_Initialize_WithAPIKey(t *testing.T) { t.Skip("E2B_API_KEY not set — skipping network-dependent init check") } p := newE2BProviderFromEnv() - ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + ctx, cancel := context.WithTimeout(t.Context(), 30*time.Second) defer cancel() if err := p.Initialize(ctx); err != nil { t.Fatalf("Initialize: %v", err) @@ -117,18 +117,19 @@ func TestE2BProvider_AllOps_BeforeInit(t *testing.T) { t.Parallel() p := newE2BProviderFromEnv() // Do NOT call Initialize. + ctx := t.Context() inst := &SandboxInstance{InstanceID: "x", Provider: ProviderE2B} - if _, err := p.CreateInstance(context.Background(), "python"); err == nil { + if _, err := p.CreateInstance(ctx, "python"); err == nil { t.Errorf("CreateInstance before init: got nil error, want one") } - if _, err := p.ExecuteCode(context.Background(), inst, "x", "python", 5, nil); err == nil { + if _, err := p.ExecuteCode(ctx, inst, "x", "python", 5, nil); err == nil { t.Errorf("ExecuteCode before init: got nil error, want one") } - if err := p.DestroyInstance(context.Background(), inst); err == nil { + if err := p.DestroyInstance(ctx, inst); err == nil { t.Errorf("DestroyInstance before init: got nil error, want one") } - if err := p.HealthCheck(context.Background()); err == nil { + if err := p.HealthCheck(ctx); err == nil { t.Errorf("HealthCheck before init: got nil error, want one") } } @@ -140,6 +141,7 @@ func TestE2BProvider_ExecuteCode_RejectsBadInputs(t *testing.T) { // — this lets us test the input-validation paths without // hitting the e2b control plane. p.initialized = true + ctx := t.Context() cases := []struct { name string @@ -149,7 +151,7 @@ func TestE2BProvider_ExecuteCode_RejectsBadInputs(t *testing.T) { { name: "empty instance id", fn: func() error { - _, err := p.ExecuteCode(context.Background(), + _, err := p.ExecuteCode(ctx, &SandboxInstance{InstanceID: ""}, "x", "python", 5, nil) return err }, @@ -158,7 +160,7 @@ func TestE2BProvider_ExecuteCode_RejectsBadInputs(t *testing.T) { { name: "nil instance", fn: func() error { - _, err := p.ExecuteCode(context.Background(), + _, err := p.ExecuteCode(ctx, nil, "x", "python", 5, nil) return err }, @@ -167,7 +169,7 @@ func TestE2BProvider_ExecuteCode_RejectsBadInputs(t *testing.T) { { name: "unsupported language", fn: func() error { - _, err := p.ExecuteCode(context.Background(), + _, err := p.ExecuteCode(ctx, &SandboxInstance{InstanceID: "x"}, "x", "ruby", 5, nil) return err }, @@ -176,7 +178,7 @@ func TestE2BProvider_ExecuteCode_RejectsBadInputs(t *testing.T) { { name: "timeout too small", fn: func() error { - _, err := p.ExecuteCode(context.Background(), + _, err := p.ExecuteCode(ctx, &SandboxInstance{InstanceID: "x"}, "x", "python", 0, nil) return err }, @@ -185,7 +187,7 @@ func TestE2BProvider_ExecuteCode_RejectsBadInputs(t *testing.T) { { name: "timeout too large", fn: func() error { - _, err := p.ExecuteCode(context.Background(), + _, err := p.ExecuteCode(ctx, &SandboxInstance{InstanceID: "x"}, "x", "python", 1000, nil) return err }, @@ -209,7 +211,8 @@ func TestE2BProvider_CreateInstance_UnsupportedLanguage(t *testing.T) { t.Parallel() p := newE2BProviderFromEnv() p.initialized = true - if _, err := p.CreateInstance(context.Background(), "ruby"); err == nil { + ctx := t.Context() + if _, err := p.CreateInstance(ctx, "ruby"); err == nil { t.Errorf("CreateInstance(ruby): got nil error, want one") } } @@ -218,10 +221,11 @@ func TestE2BProvider_DestroyInstance_EmptyID(t *testing.T) { t.Parallel() p := newE2BProviderFromEnv() p.initialized = true - if err := p.DestroyInstance(context.Background(), &SandboxInstance{InstanceID: ""}); err == nil { + ctx := t.Context() + if err := p.DestroyInstance(ctx, &SandboxInstance{InstanceID: ""}); err == nil { t.Errorf("DestroyInstance(empty id): got nil error, want one") } - if err := p.DestroyInstance(context.Background(), nil); err == nil { + if err := p.DestroyInstance(ctx, nil); err == nil { t.Errorf("DestroyInstance(nil): got nil error, want one") } } @@ -270,7 +274,7 @@ func TestE2BProvider_FullE2E_SkipWithoutKey(t *testing.T) { t.Skip("E2B_API_KEY not set — skipping full E2E test (real network call)") } p := newE2BProviderFromEnv() - ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute) + ctx, cancel := context.WithTimeout(t.Context(), 2*time.Minute) defer cancel() if err := p.Initialize(ctx); err != nil { t.Fatalf("Initialize: %v", err) @@ -318,7 +322,8 @@ func TestE2BProvider_AccessTokenFallback(t *testing.T) { // Initialize should NOT fail with "E2B_API_KEY or // E2B_ACCESS_TOKEN is required". The error we'd see is the // SDK's auth error, which is what we want. - err := p.Initialize(context.Background()) + ctx := t.Context() + err := p.Initialize(ctx) if err == nil { t.Skip("Initialize succeeded — env-var fallback accepted; skipping further checks") } diff --git a/internal/agent/sandbox/local_test.go b/internal/agent/sandbox/local_test.go index d12f1c95d8..bb896f3e43 100644 --- a/internal/agent/sandbox/local_test.go +++ b/internal/agent/sandbox/local_test.go @@ -17,7 +17,6 @@ package sandbox import ( - "context" "encoding/base64" "os" "path/filepath" @@ -43,7 +42,8 @@ func newLocalForTest(t *testing.T) *LocalProvider { maxArtifactBytes: 10 << 20, instances: map[string]string{}, } - if err := p.Initialize(context.Background()); err != nil { + ctx := t.Context() + if err := p.Initialize(ctx); err != nil { t.Fatalf("Initialize: %v", err) } return p @@ -108,7 +108,8 @@ func TestLocal_Initialize_CreatesWorkDir(t *testing.T) { } t.Setenv("LOCAL_WORK_DIR", workDir) p := newLocalProviderFromEnv() - if err := p.Initialize(context.Background()); err != nil { + ctx := t.Context() + if err := p.Initialize(ctx); err != nil { t.Fatalf("Initialize: %v", err) } info, err := os.Stat(workDir) @@ -122,7 +123,8 @@ func TestLocal_Initialize_CreatesWorkDir(t *testing.T) { func TestLocal_CreateInstance_CreatesArtifactsDir(t *testing.T) { p := newLocalForTest(t) - inst, err := p.CreateInstance(context.Background(), "python") + ctx := t.Context() + inst, err := p.CreateInstance(ctx, "python") if err != nil { t.Fatalf("CreateInstance: %v", err) } @@ -141,7 +143,8 @@ func TestLocal_CreateInstance_CreatesArtifactsDir(t *testing.T) { func TestLocal_CreateInstance_RejectsBadLanguage(t *testing.T) { p := newLocalForTest(t) - if _, err := p.CreateInstance(context.Background(), "ruby"); err == nil { + ctx := t.Context() + if _, err := p.CreateInstance(ctx, "ruby"); err == nil { t.Errorf("CreateInstance(ruby): got nil error, want one") } } @@ -149,17 +152,18 @@ func TestLocal_CreateInstance_RejectsBadLanguage(t *testing.T) { func TestLocal_AllOps_BeforeInit(t *testing.T) { t.Parallel() p := &LocalProvider{} + ctx := t.Context() inst := &SandboxInstance{InstanceID: "x", Provider: ProviderLocal} - if _, err := p.CreateInstance(context.Background(), "python"); err == nil { + if _, err := p.CreateInstance(ctx, "python"); err == nil { t.Errorf("CreateInstance before init: got nil error, want one") } - if _, err := p.ExecuteCode(context.Background(), inst, "x", "python", 5, nil); err == nil { + if _, err := p.ExecuteCode(ctx, inst, "x", "python", 5, nil); err == nil { t.Errorf("ExecuteCode before init: got nil error, want one") } - if err := p.DestroyInstance(context.Background(), inst); err == nil { + if err := p.DestroyInstance(ctx, inst); err == nil { t.Errorf("DestroyInstance before init: got nil error, want one") } - if err := p.HealthCheck(context.Background()); err == nil { + if err := p.HealthCheck(ctx); err == nil { t.Errorf("HealthCheck before init: got nil error, want one") } } @@ -173,15 +177,16 @@ func TestLocal_ExecuteCode_Python_RoundTrip(t *testing.T) { t.Skip("python3 not on PATH — skipping local subprocess test") } p := newLocalForTest(t) + ctx := t.Context() p.pythonBin = pythonPath - inst, err := p.CreateInstance(context.Background(), "python") + inst, err := p.CreateInstance(ctx, "python") if err != nil { t.Fatalf("CreateInstance: %v", err) } - defer p.DestroyInstance(context.Background(), inst) + defer p.DestroyInstance(ctx, inst) code := "def main(): return {'value': 7, 'type': 'json'}" - result, err := p.ExecuteCode(context.Background(), inst, code, "python", 10, nil) + result, err := p.ExecuteCode(ctx, inst, code, "python", 10, nil) if err != nil { t.Fatalf("ExecuteCode: %v", err) } @@ -198,7 +203,7 @@ func TestLocal_ExecuteCode_Python_RoundTrip(t *testing.T) { func TestLocal_ExecuteCode_RejectsBadInputs(t *testing.T) { p := newLocalForTest(t) p.initialized = true - + ctx := t.Context() cases := []struct { name string fn func() error @@ -207,7 +212,7 @@ func TestLocal_ExecuteCode_RejectsBadInputs(t *testing.T) { { name: "empty instance id", fn: func() error { - _, err := p.ExecuteCode(context.Background(), + _, err := p.ExecuteCode(ctx, &SandboxInstance{InstanceID: ""}, "x", "python", 5, nil) return err }, @@ -216,7 +221,7 @@ func TestLocal_ExecuteCode_RejectsBadInputs(t *testing.T) { { name: "unsupported language", fn: func() error { - _, err := p.ExecuteCode(context.Background(), + _, err := p.ExecuteCode(ctx, &SandboxInstance{InstanceID: "x"}, "x", "ruby", 5, nil) return err }, @@ -225,7 +230,7 @@ func TestLocal_ExecuteCode_RejectsBadInputs(t *testing.T) { { name: "timeout too small", fn: func() error { - _, err := p.ExecuteCode(context.Background(), + _, err := p.ExecuteCode(ctx, &SandboxInstance{InstanceID: "x"}, "x", "python", 0, nil) return err }, @@ -247,50 +252,53 @@ func TestLocal_ExecuteCode_RejectsBadInputs(t *testing.T) { func TestLocal_DestroyInstance_RemovesDir(t *testing.T) { p := newLocalForTest(t) - inst, err := p.CreateInstance(context.Background(), "python") + ctx := t.Context() + inst, err := p.CreateInstance(ctx, "python") if err != nil { t.Fatalf("CreateInstance: %v", err) } dir := filepath.Join(p.workDir, inst.InstanceID) - if _, err := os.Stat(dir); err != nil { + if _, err = os.Stat(dir); err != nil { t.Fatalf("instance dir not created: %v", err) } - if err := p.DestroyInstance(context.Background(), inst); err != nil { + if err = p.DestroyInstance(ctx, inst); err != nil { t.Errorf("DestroyInstance: %v", err) } - if _, err := os.Stat(dir); !os.IsNotExist(err) { + if _, err = os.Stat(dir); !os.IsNotExist(err) { t.Errorf("instance dir still exists after destroy: %v", err) } // Idempotent: second call should be a no-op. - if err := p.DestroyInstance(context.Background(), inst); err != nil { + if err = p.DestroyInstance(ctx, inst); err != nil { t.Errorf("DestroyInstance (idempotent): %v", err) } } func TestLocal_HealthCheck(t *testing.T) { p := newLocalForTest(t) - if err := p.HealthCheck(context.Background()); err != nil { + ctx := t.Context() + if err := p.HealthCheck(ctx); err != nil { t.Errorf("HealthCheck: %v", err) } - // Removing the work dir should make HealthCheck fail. + // Removing the work dir should make health check fail. if err := os.RemoveAll(p.workDir); err != nil { t.Fatalf("remove work dir: %v", err) } - if err := p.HealthCheck(context.Background()); err == nil { + if err := p.HealthCheck(ctx); err == nil { t.Errorf("HealthCheck after remove: got nil error, want one") } } func TestLocal_CollectArtifacts_RejectsBadExtension(t *testing.T) { p := newLocalForTest(t) - inst, err := p.CreateInstance(context.Background(), "python") + ctx := t.Context() + inst, err := p.CreateInstance(ctx, "python") if err != nil { t.Fatalf("CreateInstance: %v", err) } - defer p.DestroyInstance(context.Background(), inst) + defer p.DestroyInstance(ctx, inst) // Drop an unsupported extension into the artifacts dir. artDir := filepath.Join(p.workDir, inst.InstanceID, "artifacts") - if err := os.WriteFile(filepath.Join(artDir, "evil.exe"), []byte("x"), 0o600); err != nil { + if err = os.WriteFile(filepath.Join(artDir, "evil.exe"), []byte("x"), 0o600); err != nil { t.Fatalf("write artifact: %v", err) } _, err = p.collectArtifacts(p.workDir + "/" + inst.InstanceID) @@ -304,13 +312,14 @@ func TestLocal_CollectArtifacts_RejectsBadExtension(t *testing.T) { func TestLocal_CollectArtifacts_AllowsCSVRoundTrip(t *testing.T) { p := newLocalForTest(t) - inst, err := p.CreateInstance(context.Background(), "python") + ctx := t.Context() + inst, err := p.CreateInstance(ctx, "python") if err != nil { t.Fatalf("CreateInstance: %v", err) } - defer p.DestroyInstance(context.Background(), inst) + defer p.DestroyInstance(ctx, inst) artDir := filepath.Join(p.workDir, inst.InstanceID, "artifacts") - if err := os.WriteFile(filepath.Join(artDir, "out.csv"), []byte("a,b\n1,2\n"), 0o600); err != nil { + if err = os.WriteFile(filepath.Join(artDir, "out.csv"), []byte("a,b\n1,2\n"), 0o600); err != nil { t.Fatalf("write artifact: %v", err) } artifacts, err := p.collectArtifacts(p.workDir + "/" + inst.InstanceID) diff --git a/internal/agent/sandbox/manager_client_test.go b/internal/agent/sandbox/manager_client_test.go index ced12c97fe..b3c77c337a 100644 --- a/internal/agent/sandbox/manager_client_test.go +++ b/internal/agent/sandbox/manager_client_test.go @@ -37,7 +37,8 @@ func TestManagerClient_MapsStructuredResultToSandboxResponse(t *testing.T) { mgr.SetProvider(managerClientStubProvider{}) client := &ManagerClient{manager: mgr} - resp, err := client.ExecuteCode(context.Background(), agenttool.SandboxRequest{ + ctx := t.Context() + resp, err := client.ExecuteCode(ctx, agenttool.SandboxRequest{ Lang: "python", Script: "def main(): return 16", }) @@ -57,7 +58,8 @@ func TestManagerClient_MapsLegacyResultKeyToSandboxResponse(t *testing.T) { mgr.SetProvider(managerClientResultKeyProvider{}) client := &ManagerClient{manager: mgr} - resp, err := client.ExecuteCode(context.Background(), agenttool.SandboxRequest{ + ctx := t.Context() + resp, err := client.ExecuteCode(ctx, agenttool.SandboxRequest{ Lang: "python", Script: "def main(): return 16", }) diff --git a/internal/agent/sandbox/manager_test.go b/internal/agent/sandbox/manager_test.go index cdbf5ea01e..2f4200a9e6 100644 --- a/internal/agent/sandbox/manager_test.go +++ b/internal/agent/sandbox/manager_test.go @@ -129,8 +129,9 @@ func TestAliyun_Initialize_MissingCreds(t *testing.T) { for _, k := range []string{"AGENTRUN_ACCESS_KEY_ID", "AGENTRUN_ACCESS_KEY_SECRET", "AGENTRUN_ACCOUNT_ID"} { t.Setenv(k, "") } + ctx := t.Context() p := newAliyunProviderFromEnv() - if err := p.Initialize(context.Background()); err == nil { + if err := p.Initialize(ctx); err == nil { t.Errorf("Initialize with missing creds: got nil error, want one") } } @@ -153,19 +154,19 @@ func TestSelfManaged_EndToEnd_FullLoop(t *testing.T) { } })) defer srv.Close() - + ctx := t.Context() p := newSelfManagedForTest(srv.URL) - if err := p.Initialize(context.Background()); err != nil { + if err := p.Initialize(ctx); err != nil { t.Fatalf("Initialize: %v", err) } - inst, err := p.CreateInstance(context.Background(), "python") + inst, err := p.CreateInstance(ctx, "python") if err != nil { t.Fatalf("CreateInstance: %v", err) } if inst.Provider != ProviderSelfManaged { t.Errorf("provider = %q, want %q", inst.Provider, ProviderSelfManaged) } - result, err := p.ExecuteCode(context.Background(), inst, "def main(): return 1", "python", 5, nil) + result, err := p.ExecuteCode(ctx, inst, "def main(): return 1", "python", 5, nil) if err != nil { t.Fatalf("ExecuteCode: %v", err) } @@ -178,7 +179,7 @@ func TestSelfManaged_EndToEnd_FullLoop(t *testing.T) { if result.ExitCode != 0 { t.Errorf("exit_code = %d, want 0", result.ExitCode) } - if err := p.DestroyInstance(context.Background(), inst); err != nil { + if err = p.DestroyInstance(ctx, inst); err != nil { t.Errorf("DestroyInstance: %v", err) } } @@ -425,6 +426,7 @@ func TestLoadFromSettingsWithReader_HappyPath(t *testing.T) { w.WriteHeader(http.StatusNotFound) })) defer srv.Close() + ctx := t.Context() // Drive the mock server by setting the SANDBOX_EXECUTOR_MANAGER_URL // env var; then have the settings config return a matching @@ -443,7 +445,7 @@ func TestLoadFromSettingsWithReader_HappyPath(t *testing.T) { }, } m := &ProviderManager{} - if err := m.LoadFromSettingsWithReader(context.Background(), dao.DB, r); err != nil { + if err := m.LoadFromSettingsWithReader(ctx, dao.DB, r); err != nil { t.Fatalf("LoadFromSettingsWithReader: %v", err) } if !m.IsConfigured() { @@ -485,10 +487,11 @@ func TestLoadFromSettingsWithReader_EmptyFallback(t *testing.T) { t.Setenv("SANDBOX_PROVIDER_TYPE", "") t.Setenv("SANDBOX_EXECUTOR_MANAGER_URL", srv.URL) t.Setenv("SANDBOX_EXECUTOR_MANAGER_TIMEOUT", "5s") + ctx := t.Context() r := &fakeSettingsReader{rows: map[string][]entity.SystemSettings{}} m := &ProviderManager{} - if err := m.LoadFromSettingsWithReader(context.Background(), dao.DB, r); err != nil { + if err := m.LoadFromSettingsWithReader(ctx, dao.DB, r); err != nil { t.Fatalf("LoadFromSettingsWithReader: %v", err) } if !m.IsConfigured() { @@ -515,10 +518,11 @@ func TestLoadFromSettingsWithReader_DAOErrorFallback(t *testing.T) { t.Setenv("SANDBOX_PROVIDER_TYPE", "") t.Setenv("SANDBOX_EXECUTOR_MANAGER_URL", srv.URL) t.Setenv("SANDBOX_EXECUTOR_MANAGER_TIMEOUT", "5s") + ctx := t.Context() r := &fakeSettingsReader{fakeErr: errors.New("db is down")} m := &ProviderManager{} - if err := m.LoadFromSettingsWithReader(context.Background(), dao.DB, r); err != nil { + if err := m.LoadFromSettingsWithReader(ctx, dao.DB, r); err != nil { t.Fatalf("LoadFromSettingsWithReader (DAO error fallback): %v", err) } if got := m.Provider().ProviderType(); got != ProviderSelfManaged { @@ -542,6 +546,7 @@ func TestLoadFromSettingsWithReader_MalformedJSONFallback(t *testing.T) { t.Setenv("SANDBOX_PROVIDER_TYPE", "") t.Setenv("SANDBOX_EXECUTOR_MANAGER_URL", srv.URL) t.Setenv("SANDBOX_EXECUTOR_MANAGER_TIMEOUT", "5s") + ctx := t.Context() r := &fakeSettingsReader{ rows: map[string][]entity.SystemSettings{ @@ -550,7 +555,7 @@ func TestLoadFromSettingsWithReader_MalformedJSONFallback(t *testing.T) { }, } m := &ProviderManager{} - if err := m.LoadFromSettingsWithReader(context.Background(), dao.DB, r); err != nil { + if err := m.LoadFromSettingsWithReader(ctx, dao.DB, r); err != nil { t.Fatalf("LoadFromSettingsWithReader (malformed JSON fallback): %v", err) } sm, ok := m.Provider().(*SelfManagedProvider) @@ -582,6 +587,7 @@ func TestLoadFromSettingsWithReader_UnknownProviderType(t *testing.T) { t.Setenv("SANDBOX_PROVIDER_TYPE", "") t.Setenv("SANDBOX_EXECUTOR_MANAGER_URL", srv.URL) t.Setenv("SANDBOX_EXECUTOR_MANAGER_TIMEOUT", "5s") + ctx := t.Context() r := &fakeSettingsReader{ rows: map[string][]entity.SystemSettings{ @@ -589,7 +595,7 @@ func TestLoadFromSettingsWithReader_UnknownProviderType(t *testing.T) { }, } m := &ProviderManager{} - if err := m.LoadFromSettingsWithReader(context.Background(), dao.DB, r); err != nil { + if err := m.LoadFromSettingsWithReader(ctx, dao.DB, r); err != nil { t.Fatalf("LoadFromSettingsWithReader (unknown type fallback): %v", err) } // Falls back to env-driven self_managed, NOT the unknown type. @@ -607,13 +613,14 @@ func TestLoadFromSettingsWithReader_AlreadyLoaded_NoOp(t *testing.T) { m := &ProviderManager{} m.SetProvider(newSelfManagedProviderFromEnv()) original := m.Provider() + ctx := t.Context() r := &fakeSettingsReader{ rows: map[string][]entity.SystemSettings{ "sandbox.provider_type": {{Name: "sandbox.provider_type", Value: "local"}}, }, } - if err := m.LoadFromSettingsWithReader(context.Background(), dao.DB, r); err != nil { + if err := m.LoadFromSettingsWithReader(ctx, dao.DB, r); err != nil { t.Fatalf("LoadFromSettingsWithReader: %v", err) } if m.Provider() != original { @@ -636,6 +643,7 @@ func TestReloadFromSettingsWithReader(t *testing.T) { w.WriteHeader(http.StatusNotFound) })) defer srv.Close() + ctx := t.Context() r := &fakeSettingsReader{ rows: map[string][]entity.SystemSettings{ @@ -647,7 +655,7 @@ func TestReloadFromSettingsWithReader(t *testing.T) { }, } m := &ProviderManager{} - if err := m.ReloadFromSettingsWithReader(context.Background(), dao.DB, r); err != nil { + if err := m.ReloadFromSettingsWithReader(ctx, dao.DB, r); err != nil { t.Fatalf("ReloadFromSettingsWithReader: %v", err) } if got := m.Provider().ProviderType(); got != ProviderSelfManaged { diff --git a/internal/agent/sandbox/self_managed_test.go b/internal/agent/sandbox/self_managed_test.go index 150ed9d28a..bd0bf4bf69 100644 --- a/internal/agent/sandbox/self_managed_test.go +++ b/internal/agent/sandbox/self_managed_test.go @@ -17,7 +17,6 @@ package sandbox import ( - "context" "encoding/base64" "encoding/json" "io" @@ -50,9 +49,11 @@ func TestSelfManaged_HealthCheck_OK(t *testing.T) { w.WriteHeader(http.StatusOK) _, _ = w.Write([]byte(`{"status":"ok"}`)) })) + defer srv.Close() + ctx := t.Context() p := newSelfManagedForTest(srv.URL) - if err := p.HealthCheck(context.Background()); err != nil { + if err := p.HealthCheck(ctx); err != nil { t.Fatalf("HealthCheck: %v", err) } } @@ -63,9 +64,10 @@ func TestSelfManaged_HealthCheck_Fail(t *testing.T) { w.WriteHeader(http.StatusInternalServerError) })) defer srv.Close() + ctx := t.Context() p := newSelfManagedForTest(srv.URL) - if err := p.HealthCheck(context.Background()); err == nil { + if err := p.HealthCheck(ctx); err == nil { t.Errorf("HealthCheck on 500: got nil error, want one") } } @@ -86,9 +88,10 @@ func TestSelfManaged_Initialize(t *testing.T) { w.WriteHeader(http.StatusNotFound) })) defer srv.Close() + ctx := t.Context() p := newSelfManagedForTest(srv.URL) - if err := p.Initialize(context.Background()); err != nil { + if err := p.Initialize(ctx); err != nil { t.Fatalf("Initialize: %v", err) } if !p.isInitialized() { @@ -103,9 +106,10 @@ func TestSelfManaged_Initialize_HealthFails(t *testing.T) { w.WriteHeader(http.StatusInternalServerError) })) defer srv.Close() + ctx := t.Context() p := newSelfManagedForTest(srv.URL) - if err := p.Initialize(context.Background()); err == nil { + if err := p.Initialize(ctx); err == nil { t.Errorf("Initialize on 500 healthz: got nil error, want one") } } @@ -114,7 +118,8 @@ func TestSelfManaged_CreateInstance(t *testing.T) { t.Parallel() p := newSelfManagedForTest("http://example.invalid:9999") p.initialized = true // bypass probe for unit testing - inst, err := p.CreateInstance(context.Background(), "python") + ctx := t.Context() + inst, err := p.CreateInstance(ctx, "python") if err != nil { t.Fatalf("CreateInstance: %v", err) } @@ -133,7 +138,8 @@ func TestSelfManaged_CreateInstance_UnsupportedLanguage(t *testing.T) { t.Parallel() p := newSelfManagedForTest("http://example.invalid:9999") p.initialized = true - if _, err := p.CreateInstance(context.Background(), "ruby"); err == nil { + ctx := t.Context() + if _, err := p.CreateInstance(ctx, "ruby"); err == nil { t.Errorf("CreateInstance(ruby): got nil error, want one") } } @@ -153,14 +159,15 @@ func TestSelfManaged_ExecuteCode(t *testing.T) { }) })) defer srv.Close() + ctx := t.Context() p := newSelfManagedForTest(srv.URL) p.initialized = true - inst, err := p.CreateInstance(context.Background(), "python") + inst, err := p.CreateInstance(ctx, "python") if err != nil { t.Fatalf("CreateInstance: %v", err) } - result, err := p.ExecuteCode(context.Background(), inst, "def main(): return 1+1", "python", 10, nil) + result, err := p.ExecuteCode(ctx, inst, "def main(): return 1+1", "python", 10, nil) if err != nil { t.Fatalf("ExecuteCode: %v", err) } @@ -209,14 +216,15 @@ func TestSelfManaged_ExecuteCode_JSWrapped(t *testing.T) { handleRun(t, w, r, "ok", "") })) defer srv.Close() + ctx := t.Context() p := newSelfManagedForTest(srv.URL) p.initialized = true - inst, err := p.CreateInstance(context.Background(), "nodejs") + inst, err := p.CreateInstance(ctx, "nodejs") if err != nil { t.Fatalf("CreateInstance: %v", err) } - _, err = p.ExecuteCode(context.Background(), inst, "async function main() {}", "javascript", 5, nil) + _, err = p.ExecuteCode(ctx, inst, "async function main() {}", "javascript", 5, nil) if err != nil { t.Fatalf("ExecuteCode: %v", err) } @@ -262,14 +270,15 @@ func TestSelfManaged_ExecuteCode_PrefersHTTPResultField(t *testing.T) { }`)) })) defer srv.Close() + ctx := t.Context() p := newSelfManagedForTest(srv.URL) p.initialized = true - inst, err := p.CreateInstance(context.Background(), "python") + inst, err := p.CreateInstance(ctx, "python") if err != nil { t.Fatalf("CreateInstance: %v", err) } - result, err := p.ExecuteCode(context.Background(), inst, "def main(): return 16", "python", 10, nil) + result, err := p.ExecuteCode(ctx, inst, "def main(): return 16", "python", 10, nil) if err != nil { t.Fatalf("ExecuteCode: %v", err) } @@ -292,11 +301,12 @@ func TestSelfManaged_ExecuteCode_Non200(t *testing.T) { _, _ = w.Write([]byte("bad code")) })) defer srv.Close() + ctx := t.Context() p := newSelfManagedForTest(srv.URL) p.initialized = true - inst, _ := p.CreateInstance(context.Background(), "python") - _, err := p.ExecuteCode(context.Background(), inst, "x", "python", 5, nil) + inst, _ := p.CreateInstance(ctx, "python") + _, err := p.ExecuteCode(ctx, inst, "x", "python", 5, nil) if err == nil { t.Errorf("ExecuteCode on 400: got nil error, want one") } @@ -307,10 +317,11 @@ func TestSelfManaged_ExecuteCode_Non200(t *testing.T) { func TestSelfManaged_ExecuteCode_NotInitialized(t *testing.T) { t.Parallel() + ctx := t.Context() p := newSelfManagedForTest("http://example.invalid:9999") // do NOT set initialized inst := &SandboxInstance{InstanceID: "x"} - _, err := p.ExecuteCode(context.Background(), inst, "x", "python", 5, nil) + _, err := p.ExecuteCode(ctx, inst, "x", "python", 5, nil) if err == nil { t.Errorf("ExecuteCode on uninitialized: got nil error, want one") } @@ -318,10 +329,11 @@ func TestSelfManaged_ExecuteCode_NotInitialized(t *testing.T) { func TestSelfManaged_ExecuteCode_UnsupportedLanguage(t *testing.T) { t.Parallel() + ctx := t.Context() p := newSelfManagedForTest("http://example.invalid:9999") p.initialized = true - inst, _ := p.CreateInstance(context.Background(), "python") - _, err := p.ExecuteCode(context.Background(), inst, "x", "ruby", 5, nil) + inst, _ := p.CreateInstance(ctx, "python") + _, err := p.ExecuteCode(ctx, inst, "x", "ruby", 5, nil) if err == nil { t.Errorf("ExecuteCode(ruby): got nil error, want one") } @@ -329,9 +341,10 @@ func TestSelfManaged_ExecuteCode_UnsupportedLanguage(t *testing.T) { func TestSelfManaged_DestroyInstance_Noop(t *testing.T) { t.Parallel() + ctx := t.Context() p := newSelfManagedForTest("http://example.invalid:9999") p.initialized = true - if err := p.DestroyInstance(context.Background(), &SandboxInstance{InstanceID: "x"}); err != nil { + if err := p.DestroyInstance(ctx, &SandboxInstance{InstanceID: "x"}); err != nil { t.Errorf("DestroyInstance: %v", err) } } @@ -378,7 +391,7 @@ func TestNewSelfManagedProviderFromEnv_BaseImages(t *testing.T) { t.Errorf("nodejs baseImage = (%q, %v); want (\"\", true)", got, ok) } - // Case 3: only python set. nodejs slot must be empty. + // Case 3: only python set. Node.js slot must be empty. t.Setenv("SANDBOX_BASE_PYTHON_IMAGE", "only-python:latest") t.Setenv("SANDBOX_BASE_NODEJS_IMAGE", "") p3 := newSelfManagedProviderFromEnv() @@ -402,6 +415,7 @@ func TestSelfManaged_ExecuteCode_PassesBaseImage(t *testing.T) { handleRun(t, w, r, "ok", "") })) defer srv.Close() + ctx := t.Context() p := newSelfManagedForTest(srv.URL) p.initialized = true @@ -409,15 +423,15 @@ func TestSelfManaged_ExecuteCode_PassesBaseImage(t *testing.T) { "python": "custom-python:v1", "nodejs": "", } - inst, err := p.CreateInstance(context.Background(), "python") + inst, err := p.CreateInstance(ctx, "python") if err != nil { t.Fatalf("CreateInstance: %v", err) } - if _, err := p.ExecuteCode(context.Background(), inst, "def main(): return 1", "python", 10, nil); err != nil { + if _, err = p.ExecuteCode(ctx, inst, "def main(): return 1", "python", 10, nil); err != nil { t.Fatalf("ExecuteCode: %v", err) } var payload map[string]any - if err := json.Unmarshal(capturedBody, &payload); err != nil { + if err = json.Unmarshal(capturedBody, &payload); err != nil { t.Fatalf("decode: %v (raw=%s)", err, capturedBody) } if got := payload["base_image"]; got != "custom-python:v1" { @@ -438,6 +452,7 @@ func TestSelfManaged_ExecuteCode_OmitsEmptyBaseImage(t *testing.T) { handleRun(t, w, r, "ok", "") })) defer srv.Close() + ctx := t.Context() p := newSelfManagedForTest(srv.URL) p.initialized = true @@ -445,15 +460,15 @@ func TestSelfManaged_ExecuteCode_OmitsEmptyBaseImage(t *testing.T) { "python": "", // operator did not override "nodejs": "", } - inst, err := p.CreateInstance(context.Background(), "python") + inst, err := p.CreateInstance(ctx, "python") if err != nil { t.Fatalf("CreateInstance: %v", err) } - if _, err := p.ExecuteCode(context.Background(), inst, "def main(): return 1", "python", 10, nil); err != nil { + if _, err = p.ExecuteCode(ctx, inst, "def main(): return 1", "python", 10, nil); err != nil { t.Fatalf("ExecuteCode: %v", err) } var payload map[string]any - if err := json.Unmarshal(capturedBody, &payload); err != nil { + if err = json.Unmarshal(capturedBody, &payload); err != nil { t.Fatalf("decode: %v (raw=%s)", err, capturedBody) } if _, present := payload["base_image"]; present { diff --git a/internal/agent/sandbox/ssh_test.go b/internal/agent/sandbox/ssh_test.go index b95e4686c3..41e83cd325 100644 --- a/internal/agent/sandbox/ssh_test.go +++ b/internal/agent/sandbox/ssh_test.go @@ -102,11 +102,12 @@ func TestSSH_PrivateKeyInline(t *testing.T) { } func TestSSH_Initialize_MissingHost(t *testing.T) { + ctx := t.Context() t.Setenv("SSH_HOST", "") t.Setenv("SSH_USERNAME", "u") t.Setenv("SSH_PASSWORD", "p") p := newSSHProviderFromEnv() - if err := p.Initialize(context.Background()); err == nil { + if err := p.Initialize(ctx); err == nil { t.Errorf("Initialize with empty host: got nil error, want one") } else if !strings.Contains(err.Error(), "SSH_HOST") { t.Errorf("err = %v, want to mention SSH_HOST", err) @@ -114,23 +115,25 @@ func TestSSH_Initialize_MissingHost(t *testing.T) { } func TestSSH_Initialize_MissingUsername(t *testing.T) { + ctx := t.Context() t.Setenv("SSH_HOST", "h") t.Setenv("SSH_USERNAME", "") t.Setenv("SSH_PASSWORD", "p") p := newSSHProviderFromEnv() - if err := p.Initialize(context.Background()); err == nil { + if err := p.Initialize(ctx); err == nil { t.Errorf("Initialize with empty username: got nil error, want one") } } func TestSSH_Initialize_MissingAuth(t *testing.T) { + ctx := t.Context() t.Setenv("SSH_HOST", "h") t.Setenv("SSH_USERNAME", "u") t.Setenv("SSH_PASSWORD", "") t.Setenv("SSH_PRIVATE_KEY", "") t.Setenv("SSH_PRIVATE_KEY_PATH", "") p := newSSHProviderFromEnv() - if err := p.Initialize(context.Background()); err == nil { + if err := p.Initialize(ctx); err == nil { t.Errorf("Initialize with no auth: got nil error, want one") } else if !strings.Contains(err.Error(), "SSH_PASSWORD") { t.Errorf("err = %v, want to mention SSH_PASSWORD", err) @@ -139,31 +142,33 @@ func TestSSH_Initialize_MissingAuth(t *testing.T) { func TestSSH_AllOps_BeforeInit(t *testing.T) { t.Parallel() + ctx := t.Context() p := &SSHProvider{} inst := &SandboxInstance{InstanceID: "x", Provider: ProviderSSH} - if _, err := p.CreateInstance(context.Background(), "python"); err == nil { + if _, err := p.CreateInstance(ctx, "python"); err == nil { t.Errorf("CreateInstance before init: got nil error, want one") } - if _, err := p.ExecuteCode(context.Background(), inst, "x", "python", 5, nil); err == nil { + if _, err := p.ExecuteCode(ctx, inst, "x", "python", 5, nil); err == nil { t.Errorf("ExecuteCode before init: got nil error, want one") } - if err := p.DestroyInstance(context.Background(), inst); err == nil { + if err := p.DestroyInstance(ctx, inst); err == nil { t.Errorf("DestroyInstance before init: got nil error, want one") } - if err := p.HealthCheck(context.Background()); err == nil { + if err := p.HealthCheck(ctx); err == nil { t.Errorf("HealthCheck before init: got nil error, want one") } } func TestSSH_CreateInstance_RejectsBadLanguage(t *testing.T) { + ctx := t.Context() t.Setenv("SSH_HOST", "h") t.Setenv("SSH_USERNAME", "u") t.Setenv("SSH_PASSWORD", "p") p := newSSHProviderFromEnv() - if err := p.Initialize(context.Background()); err != nil { + if err := p.Initialize(ctx); err != nil { t.Fatalf("Initialize: %v", err) } - if _, err := p.CreateInstance(context.Background(), "ruby"); err == nil { + if _, err := p.CreateInstance(ctx, "ruby"); err == nil { t.Errorf("CreateInstance(ruby): got nil error, want one") } } @@ -172,6 +177,7 @@ func TestSSH_CreateInstance_RejectsBadLanguage(t *testing.T) { // a clear error when the host is unreachable. We bind then close // an ephemeral listener to obtain a guaranteed-closed port. func TestSSH_Dial_ConnectionRefused(t *testing.T) { + ctx := t.Context() l, err := net.Listen("tcp", "127.0.0.1:0") if err != nil { t.Fatalf("listen for ephemeral port: %v", err) @@ -185,22 +191,23 @@ func TestSSH_Dial_ConnectionRefused(t *testing.T) { t.Setenv("SSH_PASSWORD", "p") t.Setenv("SSH_TIMEOUT", "2") p := newSSHProviderFromEnv() - if err := p.Initialize(context.Background()); err != nil { + if err = p.Initialize(ctx); err != nil { t.Fatalf("Initialize: %v", err) } - ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + newCtx, cancel := context.WithTimeout(ctx, 5*time.Second) defer cancel() - if _, err := p.dial(ctx); err == nil { + if _, err = p.dial(newCtx); err == nil { t.Errorf("dial: got nil error, want one") } } func TestSSH_ExecuteCode_RejectsBadInputs(t *testing.T) { + ctx := t.Context() t.Setenv("SSH_HOST", "h") t.Setenv("SSH_USERNAME", "u") t.Setenv("SSH_PASSWORD", "p") p := newSSHProviderFromEnv() - if err := p.Initialize(context.Background()); err != nil { + if err := p.Initialize(ctx); err != nil { t.Fatalf("Initialize: %v", err) } @@ -212,7 +219,7 @@ func TestSSH_ExecuteCode_RejectsBadInputs(t *testing.T) { { name: "empty instance id", fn: func() error { - _, err := p.ExecuteCode(context.Background(), + _, err := p.ExecuteCode(ctx, &SandboxInstance{InstanceID: ""}, "x", "python", 5, nil) return err }, @@ -221,7 +228,7 @@ func TestSSH_ExecuteCode_RejectsBadInputs(t *testing.T) { { name: "unsupported language", fn: func() error { - _, err := p.ExecuteCode(context.Background(), + _, err := p.ExecuteCode(ctx, &SandboxInstance{InstanceID: "x"}, "x", "ruby", 5, nil) return err }, @@ -230,7 +237,7 @@ func TestSSH_ExecuteCode_RejectsBadInputs(t *testing.T) { { name: "unknown instance id", fn: func() error { - _, err := p.ExecuteCode(context.Background(), + _, err := p.ExecuteCode(ctx, &SandboxInstance{InstanceID: "nope"}, "x", "python", 5, nil) return err }, diff --git a/internal/agent/tool/akshare_test.go b/internal/agent/tool/akshare_test.go index 75da4799d7..a4ae975c06 100644 --- a/internal/agent/tool/akshare_test.go +++ b/internal/agent/tool/akshare_test.go @@ -67,8 +67,10 @@ func TestAkShare_FetchesStockNews(t *testing.T) { akshareStockNewsEndpoint = srv.URL + "/search/jsonp" defer func() { akshareStockNewsEndpoint = oldEndpoint }() + ctx := t.Context() + tool := NewAkShareToolWithTopN(NewHTTPHelper(), 2) - out, err := tool.InvokableRun(context.Background(), `{"query":"600519"}`) + out, err := tool.InvokableRun(ctx, `{"query":"600519"}`) if err != nil { t.Fatalf("InvokableRun: %v (out=%s)", err, out) } @@ -105,9 +107,10 @@ func TestAkShare_ParseTruncatesToTopN(t *testing.T) { func TestAkShare_RejectsMalformedJSON(t *testing.T) { t.Parallel() + ctx := t.Context() tool := NewAkShareTool() - _, err := tool.InvokableRun(context.Background(), `{not json`) + _, err := tool.InvokableRun(ctx, `{not json`) if err == nil { t.Fatal("expected error for malformed JSON, got nil") } @@ -118,9 +121,10 @@ func TestAkShare_RejectsMalformedJSON(t *testing.T) { func TestAkShare_RejectsMissingQuery(t *testing.T) { t.Parallel() + ctx := t.Context() tool := NewAkShareTool() - out, err := tool.InvokableRun(context.Background(), `{}`) + out, err := tool.InvokableRun(ctx, `{}`) if err == nil { t.Fatalf("expected error for missing query, got nil (out=%s)", out) } diff --git a/internal/agent/tool/arxiv_test.go b/internal/agent/tool/arxiv_test.go index 5b15622bb2..cee2cf40a2 100644 --- a/internal/agent/tool/arxiv_test.go +++ b/internal/agent/tool/arxiv_test.go @@ -148,9 +148,10 @@ func TestArxiv_Info(t *testing.T) { func TestArxiv_EmptyQuery(t *testing.T) { t.Parallel() + ctx := t.Context() tool := NewArxivTool() - out, err := tool.InvokableRun(context.Background(), `{"query":""}`) + out, err := tool.InvokableRun(ctx, `{"query":""}`) if err != nil { t.Fatalf("InvokableRun(empty): %v", err) } @@ -180,6 +181,7 @@ func TestArxiv_FullRoundtrip(t *testing.T) { _, _ = w.Write([]byte(canned)) })) defer srv.Close() + ctx := t.Context() // rewriteHostTransport points the hard-coded export.arxiv.org at the // test server. @@ -187,7 +189,7 @@ func TestArxiv_FullRoundtrip(t *testing.T) { Transport: rewriteHostTransport(srv.URL), }) tool := NewArxivToolWith(helper) - out, err := tool.InvokableRun(context.Background(), `{"query":"rag"}`) + out, err := tool.InvokableRun(ctx, `{"query":"rag"}`) if err != nil { t.Fatalf("InvokableRun: %v", err) } @@ -212,6 +214,7 @@ func TestArxiv_FullRoundtrip(t *testing.T) { func TestArxiv_ComponentReferencesAndDefaults(t *testing.T) { t.Parallel() + ctx := t.Context() built, err := BuildByName("arxiv", map[string]any{ "top_n": float64(7), @@ -232,7 +235,7 @@ func TestArxiv_ComponentReferencesAndDefaults(t *testing.T) { envelope := map[string]any{"results": []any{map[string]any{ "title": "Paper", "summary": "Paper summary.", "pdf_url": "https://arxiv.org/pdf/1", "entry_id": "kept", }}} - chunks, docAggs := arxiv.BuildReferences(context.Background(), envelope) + chunks, docAggs := arxiv.BuildReferences(ctx, envelope) if len(chunks) != 1 || len(docAggs) != 1 || chunks[0]["content"] != "Paper summary." { t.Fatalf("references = %#v / %#v", chunks, docAggs) } diff --git a/internal/agent/tool/code_exec_test.go b/internal/agent/tool/code_exec_test.go index c0e308e503..3028932194 100644 --- a/internal/agent/tool/code_exec_test.go +++ b/internal/agent/tool/code_exec_test.go @@ -26,9 +26,10 @@ import ( func TestCodeExec_StubsErrorWhenClientMissing(t *testing.T) { t.Parallel() + ctx := t.Context() c := NewCodeExecTool() - out, err := c.InvokableRun(context.Background(), `{"language":"python","code":"def main(): return {}"}`) + out, err := c.InvokableRun(ctx, `{"language":"python","code":"def main(): return {}"}`) if !errors.Is(err, ErrCodeExecSandboxMissing) { t.Fatalf("err = %v, want ErrCodeExecSandboxMissing", err) } @@ -47,9 +48,10 @@ func TestCodeExec_StubsErrorWhenClientMissing(t *testing.T) { func TestCodeExec_RejectsEmptyCode(t *testing.T) { t.Parallel() + ctx := t.Context() c := NewCodeExecTool() - _, err := c.InvokableRun(context.Background(), `{"language":"python","code":""}`) + _, err := c.InvokableRun(ctx, `{"language":"python","code":""}`) if err == nil || !strings.Contains(err.Error(), "code") { t.Fatalf("err = %v, want to mention empty code", err) } @@ -57,9 +59,10 @@ func TestCodeExec_RejectsEmptyCode(t *testing.T) { func TestCodeExec_RejectsBadLanguage(t *testing.T) { t.Parallel() + ctx := t.Context() c := NewCodeExecTool() - _, err := c.InvokableRun(context.Background(), `{"language":"brainfuck","code":"x"}`) + _, err := c.InvokableRun(ctx, `{"language":"brainfuck","code":"x"}`) if err == nil || !strings.Contains(err.Error(), "language") { t.Fatalf("err = %v, want to reject unsupported language", err) } @@ -67,11 +70,12 @@ func TestCodeExec_RejectsBadLanguage(t *testing.T) { func TestCodeExec_AcceptsLangAlias(t *testing.T) { t.Parallel() + ctx := t.Context() c := NewCodeExecTool() // Python tool also accepts "lang" as the field name; the Go shell // should still reach the stub branch. - _, err := c.InvokableRun(context.Background(), `{"lang":"nodejs","script":"async function main() {}"}`) + _, err := c.InvokableRun(ctx, `{"lang":"nodejs","script":"async function main() {}"}`) if !errors.Is(err, ErrCodeExecSandboxMissing) { t.Fatalf("err = %v, want ErrCodeExecSandboxMissing", err) } @@ -79,9 +83,10 @@ func TestCodeExec_AcceptsLangAlias(t *testing.T) { func TestCodeExec_Info(t *testing.T) { t.Parallel() + ctx := t.Context() c := NewCodeExecTool() - info, err := c.Info(context.Background()) + info, err := c.Info(ctx) if err != nil { t.Fatalf("Info: %v", err) } @@ -101,7 +106,7 @@ func TestCodeExec_Info(t *testing.T) { t.Fatalf("marshal Info schema: %v", err) } var schema map[string]any - if err := json.Unmarshal(encoded, &schema); err != nil { + if err = json.Unmarshal(encoded, &schema); err != nil { t.Fatalf("decode Info schema: %v", err) } properties, ok := schema["properties"].(map[string]any) @@ -109,12 +114,12 @@ func TestCodeExec_Info(t *testing.T) { t.Fatalf("Info schema properties = %#v, want object", schema["properties"]) } for _, name := range []string{"lang", "script"} { - if _, ok := properties[name]; !ok { + if _, ok = properties[name]; !ok { t.Errorf("Info schema missing %q", name) } } for _, name := range []string{"language", "code", "arguments", "outputs"} { - if _, ok := properties[name]; ok { + if _, ok = properties[name]; ok { t.Errorf("Info schema unexpectedly exposes node field %q", name) } } @@ -344,6 +349,7 @@ func TestCodeExec_ResultFallsBackToStdoutJSON(t *testing.T) { // parallel with the other CodeExec tests that depend on the // default (loud-fail) stub. func TestCodeExec_PassesTimeoutToSandbox(t *testing.T) { + ctx := t.Context() var captured SandboxRequest prev := GetSandboxClient() SetSandboxClient(stubSandbox(func(_ context.Context, req SandboxRequest) (*SandboxResponse, error) { @@ -353,7 +359,7 @@ func TestCodeExec_PassesTimeoutToSandbox(t *testing.T) { t.Cleanup(func() { SetSandboxClient(prev) }) c := NewCodeExecTool() - _, err := c.InvokableRun(context.Background(), + _, err := c.InvokableRun(ctx, `{"language":"python","code":"def main(): return {}","timeout":42}`) if err != nil { t.Fatalf("InvokableRun: %v", err) @@ -368,6 +374,7 @@ func TestCodeExec_PassesTimeoutToSandbox(t *testing.T) { // timeout test, this mutates the global sandbox client and must // not run in parallel with sibling CodeExec tests. func TestCodeExec_PassesArgumentsToSandbox(t *testing.T) { + ctx := t.Context() var captured SandboxRequest prev := GetSandboxClient() SetSandboxClient(stubSandbox(func(_ context.Context, req SandboxRequest) (*SandboxResponse, error) { @@ -377,7 +384,7 @@ func TestCodeExec_PassesArgumentsToSandbox(t *testing.T) { t.Cleanup(func() { SetSandboxClient(prev) }) c := NewCodeExecTool() - _, err := c.InvokableRun(context.Background(), + _, err := c.InvokableRun(ctx, `{"language":"python","code":"def main(**kw): return kw","arguments":{"x":1,"y":"z"}}`) if err != nil { t.Fatalf("InvokableRun: %v", err) diff --git a/internal/agent/tool/crawler_test.go b/internal/agent/tool/crawler_test.go index 4c3df14672..f263ecace8 100644 --- a/internal/agent/tool/crawler_test.go +++ b/internal/agent/tool/crawler_test.go @@ -17,7 +17,6 @@ package tool import ( - "context" "encoding/json" "errors" "net" @@ -41,6 +40,7 @@ const sampleHTML = ` func TestCrawler_FetchesAndExtractsText(t *testing.T) { t.Parallel() + ctx := t.Context() srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Type", "text/html; charset=utf-8") @@ -63,7 +63,7 @@ func TestCrawler_FetchesAndExtractsText(t *testing.T) { return host, net.ParseIP(host), nil } c := NewCrawlerTool().WithResolver(loopbackResolver) - out, err := c.InvokableRun(context.Background(), + out, err := c.InvokableRun(ctx, `{"query":`+jsonString(srv.URL)+`,"max_depth":0}`) if err != nil { t.Fatalf("InvokableRun: %v", err) @@ -105,9 +105,10 @@ func TestCrawler_FetchesAndExtractsText(t *testing.T) { func TestCrawler_RejectsMaxDepthGreaterThanZero(t *testing.T) { t.Parallel() + ctx := t.Context() c := NewCrawlerTool() - _, err := c.InvokableRun(context.Background(), `{"query":"https://example.com","max_depth":1}`) + _, err := c.InvokableRun(ctx, `{"query":"https://example.com","max_depth":1}`) if !errors.Is(err, ErrCrawlerDepthUnsupported) { t.Fatalf("err = %v, want ErrCrawlerDepthUnsupported", err) } @@ -115,9 +116,10 @@ func TestCrawler_RejectsMaxDepthGreaterThanZero(t *testing.T) { func TestCrawler_RejectsMissingQuery(t *testing.T) { t.Parallel() + ctx := t.Context() c := NewCrawlerTool() - _, err := c.InvokableRun(context.Background(), `{"query":""}`) + _, err := c.InvokableRun(ctx, `{"query":""}`) if err == nil { t.Fatal("expected error for empty query") } @@ -125,9 +127,10 @@ func TestCrawler_RejectsMissingQuery(t *testing.T) { func TestCrawler_RejectsNonHTTPScheme(t *testing.T) { t.Parallel() + ctx := t.Context() c := NewCrawlerTool() - _, err := c.InvokableRun(context.Background(), `{"query":"file:///etc/passwd"}`) + _, err := c.InvokableRun(ctx, `{"query":"file:///etc/passwd"}`) if err == nil || !strings.Contains(err.Error(), "scheme") { t.Fatalf("err = %v, want to reject file:// scheme", err) } @@ -135,6 +138,7 @@ func TestCrawler_RejectsNonHTTPScheme(t *testing.T) { func TestCrawler_AcceptsLegacyURLArgument(t *testing.T) { t.Parallel() + ctx := t.Context() sentinel := errors.New("stop after legacy url normalization") c := NewCrawlerTool().WithResolver(func(rawURL string) (string, net.IP, error) { @@ -144,7 +148,7 @@ func TestCrawler_AcceptsLegacyURLArgument(t *testing.T) { return "example.com", net.ParseIP("93.184.216.34"), sentinel }) - _, err := c.InvokableRun(context.Background(), `{"url":"https://example.com"}`) + _, err := c.InvokableRun(ctx, `{"url":"https://example.com"}`) if !errors.Is(err, sentinel) { t.Fatalf("err = %v, want resolver error after accepting legacy url", err) } @@ -152,9 +156,10 @@ func TestCrawler_AcceptsLegacyURLArgument(t *testing.T) { func TestCrawler_Info(t *testing.T) { t.Parallel() + ctx := t.Context() c := NewCrawlerTool() - info, err := c.Info(context.Background()) + info, err := c.Info(ctx) if err != nil { t.Fatalf("Info: %v", err) } diff --git a/internal/agent/tool/deepl_test.go b/internal/agent/tool/deepl_test.go index dd0c41e9d6..bf15c868ca 100644 --- a/internal/agent/tool/deepl_test.go +++ b/internal/agent/tool/deepl_test.go @@ -17,7 +17,6 @@ package tool import ( - "context" "encoding/json" "net/http" "net/http/httptest" @@ -28,6 +27,7 @@ import ( func TestDeepL_BuildRequest(t *testing.T) { t.Parallel() + ctx := t.Context() var gotMethod, gotAuth, gotCT, gotPath string var gotForm url.Values @@ -52,7 +52,7 @@ func TestDeepL_BuildRequest(t *testing.T) { Transport: rewriteHostTransport(srv.URL), }) tool := NewDeepLToolWith(helper) - out, err := tool.InvokableRun(context.Background(), + out, err := tool.InvokableRun(ctx, `{"api_key":"key-xyz:fx","text":"Hello world","source_lang":"en","target_lang":"de"}`) if err != nil { t.Fatalf("InvokableRun: %v", err) @@ -98,6 +98,7 @@ func TestDeepL_BuildRequest(t *testing.T) { func TestDeepL_DefaultLanguages(t *testing.T) { t.Parallel() + ctx := t.Context() var gotForm url.Values srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { @@ -112,7 +113,7 @@ func TestDeepL_DefaultLanguages(t *testing.T) { Transport: rewriteHostTransport(srv.URL), }) tool := NewDeepLToolWith(helper) - if _, err := tool.InvokableRun(context.Background(), + if _, err := tool.InvokableRun(ctx, `{"api_key":"x:fx","text":"Hello"}`); err != nil { t.Fatalf("InvokableRun: %v", err) } @@ -126,13 +127,14 @@ func TestDeepL_DefaultLanguages(t *testing.T) { func TestDeepL_RequiresAPIKeyAndText(t *testing.T) { t.Parallel() + ctx := t.Context() tool := NewDeepLTool() - if _, err := tool.InvokableRun(context.Background(), + if _, err := tool.InvokableRun(ctx, `{"api_key":"","text":"Hello"}`); err == nil { t.Error("expected error for missing api_key") } - if _, err := tool.InvokableRun(context.Background(), + if _, err := tool.InvokableRun(ctx, `{"api_key":"x","text":""}`); err == nil { t.Error("expected error for empty text") } @@ -140,9 +142,10 @@ func TestDeepL_RequiresAPIKeyAndText(t *testing.T) { func TestDeepL_Info(t *testing.T) { t.Parallel() + ctx := t.Context() tool := NewDeepLTool() - info, err := tool.Info(context.Background()) + info, err := tool.Info(ctx) if err != nil { t.Fatalf("Info: %v", err) } @@ -178,6 +181,7 @@ func TestDeepL_Info(t *testing.T) { // when both tests run in the same package. func TestDeepL_TranslationFailureReturnsError(t *testing.T) { t.Parallel() + ctx := t.Context() // 500 Internal Server Error from a stub DeepL endpoint. srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { @@ -189,7 +193,7 @@ func TestDeepL_TranslationFailureReturnsError(t *testing.T) { Transport: rewriteHostTransport(srv.URL), }) tool := NewDeepLToolWith(helper) - out, err := tool.InvokableRun(context.Background(), + out, err := tool.InvokableRun(ctx, `{"api_key":"key-xyz:fx","text":"hello","source_lang":"EN","target_lang":"ZH"}`) if err == nil { t.Fatalf("expected non-nil error, got nil; out=%s", out) diff --git a/internal/agent/tool/duckduckgo_test.go b/internal/agent/tool/duckduckgo_test.go index 435a80158a..a930e2266b 100644 --- a/internal/agent/tool/duckduckgo_test.go +++ b/internal/agent/tool/duckduckgo_test.go @@ -17,7 +17,6 @@ package tool import ( - "context" "encoding/json" "net/http" "net/http/httptest" @@ -92,6 +91,7 @@ func TestDuckDuckGo_BuildNewsURLWithVQD(t *testing.T) { } func TestDuckDuckGo_ParseGeneralResults(t *testing.T) { + ctx := t.Context() srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Type", "text/html; charset=utf-8") _, _ = w.Write([]byte(` @@ -114,7 +114,7 @@ func TestDuckDuckGo_ParseGeneralResults(t *testing.T) { t.Cleanup(func() { duckduckgoSearchEndpoint = prevSearch }) tool := NewDuckDuckGoTool() - out, err := tool.InvokableRun(context.Background(), `{"query":"ragflow","top_n":5}`) + out, err := tool.InvokableRun(ctx, `{"query":"ragflow","top_n":5}`) if err != nil { t.Fatalf("InvokableRun: %v", err) } @@ -141,6 +141,7 @@ func TestDuckDuckGo_ParseGeneralResults(t *testing.T) { } func TestDuckDuckGo_ParseNewsResults(t *testing.T) { + ctx := t.Context() srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { switch r.URL.Path { case "/bootstrap": @@ -171,7 +172,7 @@ func TestDuckDuckGo_ParseNewsResults(t *testing.T) { t.Cleanup(func() { duckduckgoNewsBootstrapEndpoint = prevBootstrap }) tool := NewDuckDuckGoTool() - out, err := tool.InvokableRun(context.Background(), `{"query":"ragflow","channel":"news","top_n":1}`) + out, err := tool.InvokableRun(ctx, `{"query":"ragflow","channel":"news","top_n":1}`) if err != nil { t.Fatalf("InvokableRun: %v", err) } @@ -195,6 +196,7 @@ func TestDuckDuckGo_ParseNewsResults(t *testing.T) { } func TestDuckDuckGo_DefaultChannelUsesGeneralSearch(t *testing.T) { + ctx := t.Context() srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if got := r.URL.Path; got != "/" { // keep old behavior impossible to hit if search endpoint override works incorrectly @@ -214,15 +216,16 @@ func TestDuckDuckGo_DefaultChannelUsesGeneralSearch(t *testing.T) { t.Cleanup(func() { duckduckgoSearchEndpoint = prevSearch }) tool := NewDuckDuckGoTool() - _, err := tool.InvokableRun(context.Background(), `{"query":"ragflow"}`) + _, err := tool.InvokableRun(ctx, `{"query":"ragflow"}`) if err != nil { t.Fatalf("InvokableRun: %v", err) } } func TestDuckDuckGo_Info(t *testing.T) { + ctx := t.Context() tool := NewDuckDuckGoTool() - info, err := tool.Info(context.Background()) + info, err := tool.Info(ctx) if err != nil { t.Fatalf("Info: %v", err) } @@ -253,18 +256,20 @@ func TestDuckDuckGo_Info(t *testing.T) { } func TestDuckDuckGo_EmptyQuery(t *testing.T) { + ctx := t.Context() tool := NewDuckDuckGoTool() - out, err := tool.InvokableRun(context.Background(), `{"query":""}`) + out, err := tool.InvokableRun(ctx, `{"query":""}`) if err != nil { t.Fatalf("InvokableRun(empty): %v", err) } var envelope duckduckgoEnvelope - if err := json.Unmarshal([]byte(out), &envelope); err != nil || len(envelope.Results) != 0 { + if err = json.Unmarshal([]byte(out), &envelope); err != nil || len(envelope.Results) != 0 { t.Fatalf("empty result = %s / %v", out, err) } } func TestDuckDuckGo_RealReactAgent_ExecutesTool(t *testing.T) { + ctx := t.Context() var hitCount int srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { hitCount++ @@ -290,7 +295,7 @@ func TestDuckDuckGo_RealReactAgent_ExecutesTool(t *testing.T) { "RAGFlow is an open-source RAG engine.", ) - agent, err := react.NewAgent(context.Background(), &react.AgentConfig{ + agent, err := react.NewAgent(ctx, &react.AgentConfig{ ToolCallingModel: mdl, ToolsConfig: compose.ToolsNodeConfig{ Tools: []einotool.BaseTool{realTool}, @@ -301,7 +306,7 @@ func TestDuckDuckGo_RealReactAgent_ExecutesTool(t *testing.T) { t.Fatalf("react.NewAgent: %v", err) } - out, err := agent.Generate(context.Background(), []*schema.Message{ + out, err := agent.Generate(ctx, []*schema.Message{ schema.UserMessage("What is RAGFlow?"), }) if err != nil { @@ -339,6 +344,7 @@ func TestDuckDuckGo_RealReactAgent_ExecutesTool(t *testing.T) { } func TestDuckDuckGo_ComponentReferencesAndDefaults(t *testing.T) { + ctx := t.Context() t.Parallel() built, err := BuildByName("duckduckgo", map[string]any{ @@ -365,7 +371,7 @@ func TestDuckDuckGo_ComponentReferencesAndDefaults(t *testing.T) { envelope := map[string]any{"results": []any{map[string]any{ "title": "Story", "url": "https://news.example/story", "body": "Breaking update", }}} - chunks, docAggs := duck.BuildReferences(context.Background(), envelope) + chunks, docAggs := duck.BuildReferences(ctx, envelope) if len(chunks) != 1 || len(docAggs) != 1 || chunks[0]["content"] != "Breaking update" { t.Fatalf("references = %#v / %#v", chunks, docAggs) } diff --git a/internal/agent/tool/email_test.go b/internal/agent/tool/email_test.go index 1e2142d395..f3d2f149af 100644 --- a/internal/agent/tool/email_test.go +++ b/internal/agent/tool/email_test.go @@ -64,6 +64,7 @@ func TestEmail_BuildMessage(t *testing.T) { } func TestEmail_SendBuildsDistinctHeadersAndEnvelopeRecipients(t *testing.T) { + ctx := t.Context() originalSendEmail := sendEmail t.Cleanup(func() { sendEmail = originalSendEmail }) var sentParams emailParams @@ -92,12 +93,12 @@ func TestEmail_SendBuildsDistinctHeadersAndEnvelopeRecipients(t *testing.T) { argsJSON, _ := json.Marshal(args) state := runtime.NewCanvasState("run-email", "task-email") state.Sys["date"] = "2026-07-15" - out, err := built.(*EmailTool).InvokableRun(runtime.WithState(context.Background(), state), string(argsJSON)) + out, err := built.(*EmailTool).InvokableRun(runtime.WithState(ctx, state), string(argsJSON)) if err != nil { t.Fatalf("InvokableRun: %v", err) } var env emailEnvelope - if err := json.Unmarshal([]byte(out), &env); err != nil || !env.OK || env.Error != "" { + if err = json.Unmarshal([]byte(out), &env); err != nil || !env.OK || env.Error != "" { t.Fatalf("output = %s, decode error = %v", out, err) } @@ -119,6 +120,7 @@ func TestEmail_SendBuildsDistinctHeadersAndEnvelopeRecipients(t *testing.T) { } func TestEmail_STARTTLSRequiredBeforeSubmission(t *testing.T) { + ctx := t.Context() t.Parallel() listener, err := net.Listen("tcp", "127.0.0.1:0") @@ -161,7 +163,7 @@ func TestEmail_STARTTLSRequiredBeforeSubmission(t *testing.T) { } var portNumber int _, _ = fmt.Sscanf(port, "%d", &portNumber) - err = sendEmailSTARTTLS(context.Background(), emailParams{ + err = sendEmailSTARTTLS(ctx, emailParams{ SMTPServer: host, SMTPPort: portNumber, Email: "alice@example.com", ToEmail: "bob@example.com", }, []byte("message")) @@ -182,6 +184,7 @@ func TestEmail_STARTTLSRequiredBeforeSubmission(t *testing.T) { } func TestEmail_RequiresFields(t *testing.T) { + ctx := t.Context() t.Parallel() cases := []struct { @@ -212,7 +215,7 @@ func TestEmail_RequiresFields(t *testing.T) { for _, tc := range cases { t.Run(tc.name, func(t *testing.T) { t.Parallel() - _, err := tc.tool.InvokableRun(context.Background(), tc.args) + _, err := tc.tool.InvokableRun(ctx, tc.args) if err == nil { t.Fatalf("expected error for %s", tc.name) } @@ -224,10 +227,11 @@ func TestEmail_RequiresFields(t *testing.T) { } func TestEmail_Info(t *testing.T) { + ctx := t.Context() t.Parallel() tool := NewEmailTool() - info, err := tool.Info(context.Background()) + info, err := tool.Info(ctx) if err != nil { t.Fatalf("Info: %v", err) } diff --git a/internal/agent/tool/exesql_test.go b/internal/agent/tool/exesql_test.go index c3f5f7ce0e..6438340ac1 100644 --- a/internal/agent/tool/exesql_test.go +++ b/internal/agent/tool/exesql_test.go @@ -65,6 +65,7 @@ func sqlmockDialer(t *testing.T) (exesqlDialer, sqlmock.Sqlmock, func()) { } func TestExeSQL_NoCredentials(t *testing.T) { + ctx := t.Context() t.Parallel() e := NewExeSQLTool(exesqlConnParams{}). @@ -72,7 +73,7 @@ func TestExeSQL_NoCredentials(t *testing.T) { t.Fatal("dialer should not be called when credentials are missing") return nil, nil }) - _, err := e.InvokableRun(context.Background(), `{"sql":"SELECT 1"}`) + _, err := e.InvokableRun(ctx, `{"sql":"SELECT 1"}`) if !errors.Is(err, ErrExeSQLNoCredentials) { t.Fatalf("err = %v, want ErrExeSQLNoCredentials", err) } @@ -108,6 +109,7 @@ func TestExeSQL_RejectsNonSelect(t *testing.T) { {"merge cte", `WITH changed AS (MERGE INTO users USING incoming ON users.id = incoming.id WHEN MATCHED THEN UPDATE SET name = incoming.name RETURNING *) SELECT * FROM changed`}, } + ctx := t.Context() for _, c := range cases { t.Run(c.name, func(t *testing.T) { t.Parallel() @@ -116,7 +118,7 @@ func TestExeSQL_RejectsNonSelect(t *testing.T) { t.Fatal("dialer called for rejected SQL") return nil, nil }) - _, err := e.InvokableRun(context.Background(), + _, err := e.InvokableRun(ctx, `{"sql":`+jsonString(c.sql)+`}`) if !errors.Is(err, ErrExeSQLNotSelect) { t.Fatalf("err = %v, want ErrExeSQLNotSelect", err) @@ -131,8 +133,9 @@ func TestExeSQL_RejectsMixedBatchBeforeDatabaseAccess(t *testing.T) { t.Fatal("dialer called before every SQL statement was validated") return nil, nil }) + ctx := t.Context() - _, err := e.InvokableRun(context.Background(), `{"sql":"SELECT 1; DROP TABLE users"}`) + _, err := e.InvokableRun(ctx, `{"sql":"SELECT 1; DROP TABLE users"}`) if !errors.Is(err, ErrExeSQLNotSelect) { t.Fatalf("err = %v, want ErrExeSQLNotSelect", err) } @@ -158,6 +161,7 @@ func TestExeSQL_AllowsSelect(t *testing.T) { // Block comment. `/* DROP TABLE foo */ SELECT 1`, } + ctx := t.Context() for _, sql := range cases { t.Run(sql, func(t *testing.T) { t.Parallel() @@ -172,7 +176,7 @@ func TestExeSQL_AllowsSelect(t *testing.T) { sqlmock.NewRows([]string{"1"}), ) e := NewExeSQLTool(testConn()).WithExeSQLDialer(dialer) - _, err := e.InvokableRun(context.Background(), + _, err := e.InvokableRun(ctx, `{"sql":`+jsonString(sql)+`}`) // Two acceptable outcomes: // 1. SQL is the literal `SELECT 1` and matches the @@ -196,20 +200,22 @@ func TestExeSQL_AllowsSelect(t *testing.T) { } func TestExeSQL_RejectsEmptySQL(t *testing.T) { + ctx := t.Context() t.Parallel() e := NewExeSQLTool(testConn()) - _, err := e.InvokableRun(context.Background(), `{"sql":""}`) + _, err := e.InvokableRun(ctx, `{"sql":""}`) if err == nil || !strings.Contains(err.Error(), "sql") { t.Fatalf("err = %v, want to mention empty sql", err) } } func TestExeSQL_RejectsEmptyArgs(t *testing.T) { + ctx := t.Context() t.Parallel() e := NewExeSQLTool(testConn()) - _, err := e.InvokableRun(context.Background(), "") + _, err := e.InvokableRun(ctx, "") if err == nil { t.Fatal("expected error for empty args") } @@ -272,6 +278,7 @@ func TestExeSQL_ReadOnlyValidationIgnoresQuotedAndCommentedKeywords(t *testing.T } func TestExeSQL_ExecutesStatementsWithQuotedSemicolonsIntact(t *testing.T) { + ctx := t.Context() t.Parallel() dialer, mock, cleanup := sqlmockDialer(t) @@ -283,7 +290,7 @@ func TestExeSQL_ExecutesStatementsWithQuotedSemicolonsIntact(t *testing.T) { WillReturnRows(sqlmock.NewRows([]string{"value"}).AddRow(2)) e := NewExeSQLTool(testConn()).WithExeSQLDialer(dialer) - if _, err := e.InvokableRun(context.Background(), `{"sql":"SELECT 'hello; world'; SELECT 2"}`); err != nil { + if _, err := e.InvokableRun(ctx, `{"sql":"SELECT 'hello; world'; SELECT 2"}`); err != nil { t.Fatalf("InvokableRun: %v", err) } if err := mock.ExpectationsWereMet(); err != nil { @@ -292,6 +299,7 @@ func TestExeSQL_ExecutesStatementsWithQuotedSemicolonsIntact(t *testing.T) { } func TestExeSQL_RejectsMySQLExecutableComment(t *testing.T) { + ctx := t.Context() t.Parallel() e := NewExeSQLTool(testConn()). @@ -299,17 +307,18 @@ func TestExeSQL_RejectsMySQLExecutableComment(t *testing.T) { t.Fatal("dialer called for an executable comment") return nil, nil }) - _, err := e.InvokableRun(context.Background(), `{"sql":"SELECT 1 /*!; DROP TABLE users */"}`) + _, err := e.InvokableRun(ctx, `{"sql":"SELECT 1 /*!; DROP TABLE users */"}`) if !errors.Is(err, ErrExeSQLNotSelect) { t.Fatalf("err = %v, want ErrExeSQLNotSelect", err) } } func TestExeSQL_Info(t *testing.T) { + ctx := t.Context() t.Parallel() e := NewExeSQLTool(testConn()) - info, err := e.Info(context.Background()) + info, err := e.Info(ctx) if err != nil { t.Fatalf("Info: %v", err) } @@ -337,6 +346,7 @@ func TestExeSQL_Info(t *testing.T) { } func TestExeSQL_UsesConfiguredSQLDefault(t *testing.T) { + ctx := t.Context() dialer, mock, cleanup := sqlmockDialer(t) defer cleanup() mock.ExpectPing() @@ -347,7 +357,7 @@ func TestExeSQL_UsesConfiguredSQLDefault(t *testing.T) { conn.SQL = "SELECT 1" e := NewExeSQLTool(conn).WithExeSQLDialer(dialer) - out, err := e.InvokableRun(context.Background(), `{}`) + out, err := e.InvokableRun(ctx, `{}`) if err != nil { t.Fatalf("InvokableRun: %v", err) } @@ -360,6 +370,7 @@ func TestExeSQL_UsesConfiguredSQLDefault(t *testing.T) { } func TestExeSQL_ComponentContractAndTemplateResolution(t *testing.T) { + ctx := t.Context() dialer, mock, cleanup := sqlmockDialer(t) defer cleanup() mock.ExpectPing() @@ -370,12 +381,12 @@ func TestExeSQL_ComponentContractAndTemplateResolution(t *testing.T) { exesql := NewExeSQLTool(conn).WithExeSQLDialer(dialer) state := runtime.NewCanvasState("run", "task") state.SetVar("Agent:Result", "content", "SELECT id FROM orders WHERE status = 'Completed'") - out, err := exesql.InvokableRun(runtime.WithState(context.Background(), state), `{}`) + out, err := exesql.InvokableRun(runtime.WithState(ctx, state), `{}`) if err != nil { t.Fatalf("InvokableRun: %v", err) } var envelope map[string]any - if err := json.Unmarshal([]byte(out), &envelope); err != nil { + if err = json.Unmarshal([]byte(out), &envelope); err != nil { t.Fatalf("decode output: %v", err) } outputs := exesql.BuildComponentOutputs(envelope) @@ -390,7 +401,7 @@ func TestExeSQL_ComponentContractAndTemplateResolution(t *testing.T) { if sqlInput, ok := spec.InputForm["sql"].(map[string]any); !ok || sqlInput["type"] != "line" { t.Fatalf("sql input form = %#v", spec.InputForm["sql"]) } - if err := mock.ExpectationsWereMet(); err != nil { + if err = mock.ExpectationsWereMet(); err != nil { t.Fatalf("sql expectations: %v", err) } } @@ -410,6 +421,7 @@ func TestExeSQL_BuildByNameAcceptsCanvasShape(t *testing.T) { } func TestExeSQL_ExecuteSelect_ReturnsRows(t *testing.T) { + ctx := t.Context() t.Parallel() dialer, mock, cleanup := sqlmockDialer(t) @@ -423,7 +435,7 @@ func TestExeSQL_ExecuteSelect_ReturnsRows(t *testing.T) { AddRow(8, "bob")) e := NewExeSQLTool(testConn()).WithExeSQLDialer(dialer) - out, err := e.InvokableRun(context.Background(), + out, err := e.InvokableRun(ctx, `{"sql":"SELECT id, name FROM t WHERE id = 7"}`) if err != nil { t.Fatalf("InvokableRun: %v", err) @@ -444,6 +456,7 @@ func TestExeSQL_ExecuteSelect_ReturnsRows(t *testing.T) { } func TestExeSQL_ExecuteSelect_NoRowsReturnsSentinel(t *testing.T) { + ctx := t.Context() t.Parallel() dialer, mock, cleanup := sqlmockDialer(t) @@ -453,12 +466,12 @@ func TestExeSQL_ExecuteSelect_NoRowsReturnsSentinel(t *testing.T) { WillReturnRows(sqlmock.NewRows([]string{"x"})) e := NewExeSQLTool(testConn()).WithExeSQLDialer(dialer) - out, err := e.InvokableRun(context.Background(), `{"sql":"SELECT 1"}`) + out, err := e.InvokableRun(ctx, `{"sql":"SELECT 1"}`) if err != nil { t.Fatalf("InvokableRun: %v", err) } var got exesqlResult - if err := json.Unmarshal([]byte(out), &got); err != nil { + if err = json.Unmarshal([]byte(out), &got); err != nil { t.Fatalf("unmarshal: %v\nout=%s", err, out) } // The Python tool's "No record in the database!" sentinel must @@ -470,6 +483,7 @@ func TestExeSQL_ExecuteSelect_NoRowsReturnsSentinel(t *testing.T) { } func TestExeSQL_ExecuteSelect_PerStatementErrorIsolated(t *testing.T) { + ctx := t.Context() t.Parallel() dialer, mock, cleanup := sqlmockDialer(t) @@ -484,13 +498,13 @@ func TestExeSQL_ExecuteSelect_PerStatementErrorIsolated(t *testing.T) { WillReturnError(errors.New("syntax error at or near BOGUS")) e := NewExeSQLTool(testConn()).WithExeSQLDialer(dialer) - out, err := e.InvokableRun(context.Background(), + out, err := e.InvokableRun(ctx, `{"sql":"SELECT 1; SELECT * FROM bogus"}`) if err != nil { t.Fatalf("InvokableRun should not abort on a per-statement error: %v", err) } var got exesqlResult - if err := json.Unmarshal([]byte(out), &got); err != nil { + if err = json.Unmarshal([]byte(out), &got); err != nil { t.Fatalf("unmarshal: %v\nout=%s", err, out) } if len(got.Rows) != 2 { @@ -507,6 +521,7 @@ func TestExeSQL_ExecuteSelect_PerStatementErrorIsolated(t *testing.T) { } func TestExeSQL_ExecuteSelect_NormalizesTimeAndBytes(t *testing.T) { + ctx := t.Context() t.Parallel() dialer, mock, cleanup := sqlmockDialer(t) @@ -517,13 +532,13 @@ func TestExeSQL_ExecuteSelect_NormalizesTimeAndBytes(t *testing.T) { AddRow("2024-06-12T03:04:05Z", []byte("hello"))) e := NewExeSQLTool(testConn()).WithExeSQLDialer(dialer) - out, err := e.InvokableRun(context.Background(), + out, err := e.InvokableRun(ctx, `{"sql":"SELECT ts, blob_col FROM t"}`) if err != nil { t.Fatalf("InvokableRun: %v", err) } var got exesqlResult - if err := json.Unmarshal([]byte(out), &got); err != nil { + if err = json.Unmarshal([]byte(out), &got); err != nil { t.Fatalf("unmarshal: %v\nout=%s", err, out) } if len(got.Rows) != 1 { @@ -540,6 +555,7 @@ func TestExeSQL_ExecuteSelect_NormalizesTimeAndBytes(t *testing.T) { } func TestExeSQL_UnsupportedDB(t *testing.T) { + ctx := t.Context() t.Parallel() e := NewExeSQLTool(exesqlConnParams{ @@ -547,7 +563,7 @@ func TestExeSQL_UnsupportedDB(t *testing.T) { Host: "1.1.1.1", Port: 8080, Database: "catalog", Username: "u", Password: "p", }) - _, err := e.InvokableRun(context.Background(), `{"sql":"SELECT 1"}`) + _, err := e.InvokableRun(ctx, `{"sql":"SELECT 1"}`) if err == nil { t.Fatal("expected non-nil error for trino without registered driver") } @@ -747,6 +763,7 @@ func (m *reactScriptedModel) Stream(_ context.Context, _ []*schema.Message, _ .. // and the resulting JSON is passed back as a ToolMessage. Replacing // the model with a hand-rolled stub would skip all of that. func TestExeSQL_RealReactAgent_ExecutesTool(t *testing.T) { + ctx := t.Context() t.Parallel() dialer, mock, cleanup := sqlmockDialer(t) @@ -763,7 +780,7 @@ func TestExeSQL_RealReactAgent_ExecutesTool(t *testing.T) { "the answer is 42", ) - agent, err := react.NewAgent(context.Background(), &react.AgentConfig{ + agent, err := react.NewAgent(ctx, &react.AgentConfig{ ToolCallingModel: mdl, ToolsConfig: compose.ToolsNodeConfig{ Tools: []einotool.BaseTool{realTool}, @@ -774,7 +791,7 @@ func TestExeSQL_RealReactAgent_ExecutesTool(t *testing.T) { t.Fatalf("react.NewAgent: %v", err) } - out, err := agent.Generate(context.Background(), []*schema.Message{ + out, err := agent.Generate(ctx, []*schema.Message{ schema.UserMessage("What is 42?"), }) if err != nil { @@ -806,7 +823,7 @@ func TestExeSQL_RealReactAgent_ExecutesTool(t *testing.T) { if !sawToolResult { t.Errorf("round 2 input did not contain a ToolMessage carrying the tool result; got %d messages", len(mdl.rounds[1])) } - if err := mock.ExpectationsWereMet(); err != nil { + if err = mock.ExpectationsWereMet(); err != nil { t.Errorf("sqlmock expectations not met: %v", err) } } @@ -818,6 +835,7 @@ func TestExeSQL_RealReactAgent_ExecutesTool(t *testing.T) { // the model on round 2 without crashing the ReAct loop, so the model // can ground its final answer in the surfaced error. func TestExeSQL_RealReactAgent_ToolErrorIsolated(t *testing.T) { + ctx := t.Context() t.Parallel() dialer, mock, cleanup := sqlmockDialer(t) @@ -834,7 +852,7 @@ func TestExeSQL_RealReactAgent_ToolErrorIsolated(t *testing.T) { "the query had a syntax error", ) - agent, err := react.NewAgent(context.Background(), &react.AgentConfig{ + agent, err := react.NewAgent(ctx, &react.AgentConfig{ ToolCallingModel: mdl, ToolsConfig: compose.ToolsNodeConfig{ Tools: []einotool.BaseTool{realTool}, @@ -845,7 +863,7 @@ func TestExeSQL_RealReactAgent_ToolErrorIsolated(t *testing.T) { t.Fatalf("react.NewAgent: %v", err) } - out, err := agent.Generate(context.Background(), []*schema.Message{ + out, err := agent.Generate(ctx, []*schema.Message{ schema.UserMessage("Find bogus rows"), }) if err != nil { @@ -869,7 +887,7 @@ func TestExeSQL_RealReactAgent_ToolErrorIsolated(t *testing.T) { if !sawErrorResult { t.Errorf("round 2 input did not contain a ToolMessage with the DB error; got %d messages", len(mdl.rounds[1])) } - if err := mock.ExpectationsWereMet(); err != nil { + if err = mock.ExpectationsWereMet(); err != nil { t.Errorf("sqlmock expectations: %v", err) } } diff --git a/internal/agent/tool/exesql_trino_test.go b/internal/agent/tool/exesql_trino_test.go index d613899240..17502514ae 100644 --- a/internal/agent/tool/exesql_trino_test.go +++ b/internal/agent/tool/exesql_trino_test.go @@ -17,7 +17,6 @@ package tool import ( - "context" "strings" "testing" @@ -208,7 +207,8 @@ func TestExeSQL_Trino_HappyPath(t *testing.T) { MaxRecords: 100, }).WithExeSQLDialer(dialer) - out, err := tool.InvokableRun(context.Background(), + ctx := t.Context() + out, err := tool.InvokableRun(ctx, `{"sql":"SELECT id, name FROM catalog.tiny.users"}`) if err != nil { t.Fatalf("InvokableRun: %v", err) diff --git a/internal/agent/tool/exesql_unsupported_test.go b/internal/agent/tool/exesql_unsupported_test.go index 0437c97e4d..4eb29971b4 100644 --- a/internal/agent/tool/exesql_unsupported_test.go +++ b/internal/agent/tool/exesql_unsupported_test.go @@ -17,7 +17,6 @@ package tool import ( - "context" "errors" "testing" ) @@ -27,9 +26,10 @@ import ( // database/sql driver, so InvokableRun should fail at sql.Open with an // unknown-driver error rather than the old unsupported-db sentinel. func TestExeSQL_TrinoDriverMissing(t *testing.T) { + ctx := t.Context() conn := exesqlConnParams{DBType: "trino", Host: "1.1.1.1", Port: 8080, Database: "d", Username: "u"} tool := NewExeSQLTool(conn) - _, err := tool.InvokableRun(context.Background(), `{"sql":"SELECT 1"}`) + _, err := tool.InvokableRun(ctx, `{"sql":"SELECT 1"}`) if err == nil { t.Fatal("expected driver error for trino") } @@ -40,9 +40,10 @@ func TestExeSQL_TrinoDriverMissing(t *testing.T) { // TestExeSQL_IBMDB2Unsupported: same as above for IBM DB2. func TestExeSQL_IBMDB2Unsupported(t *testing.T) { + ctx := t.Context() conn := exesqlConnParams{DBType: "ibm db2", Host: "1.1.1.1", Port: 50000, Database: "d", Username: "u"} tool := NewExeSQLTool(conn) - _, err := tool.InvokableRun(context.Background(), `{"sql":"SELECT 1"}`) + _, err := tool.InvokableRun(ctx, `{"sql":"SELECT 1"}`) if err == nil { t.Fatal("expected ErrExeSQLUnsupportedDB for ibm db2") } @@ -57,9 +58,10 @@ func TestExeSQL_IBMDB2Unsupported(t *testing.T) { // follow-up should normalize the error. The regression guard here // is "doesn't panic, returns a non-nil error". func TestExeSQL_UnknownDB(t *testing.T) { + ctx := t.Context() conn := exesqlConnParams{DBType: "fake-db", Host: "1.1.1.1", Port: 1234, Database: "d", Username: "u"} tool := NewExeSQLTool(conn) - _, err := tool.InvokableRun(context.Background(), `{"sql":"SELECT 1"}`) + _, err := tool.InvokableRun(ctx, `{"sql":"SELECT 1"}`) if err == nil { t.Fatal("expected error for unknown db_type") } diff --git a/internal/agent/tool/github_test.go b/internal/agent/tool/github_test.go index 281e5c6f5f..338695cd66 100644 --- a/internal/agent/tool/github_test.go +++ b/internal/agent/tool/github_test.go @@ -17,7 +17,6 @@ package tool import ( - "context" "encoding/json" "net/http" "net/http/httptest" @@ -91,6 +90,7 @@ func TestGitHub_BuildURL(t *testing.T) { func TestGitHub_ParseResponse(t *testing.T) { t.Parallel() + ctx := t.Context() var gotContentType, gotAPIVersion, gotPerPage string srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { @@ -111,7 +111,7 @@ func TestGitHub_ParseResponse(t *testing.T) { Transport: rewriteHostTransport(srv.URL), }) tool := NewGitHubToolWithDefaults(helper, githubParams{TopN: 17}) - out, err := tool.InvokableRun(context.Background(), + out, err := tool.InvokableRun(ctx, `{"query":"ragflow"}`) if err != nil { t.Fatalf("InvokableRun: %v", err) @@ -149,14 +149,15 @@ func TestGitHub_ParseResponse(t *testing.T) { func TestGitHub_EmptyQueryReturnsEmptyResults(t *testing.T) { t.Parallel() + ctx := t.Context() tool := NewGitHubTool() - out, err := tool.InvokableRun(context.Background(), `{"query":""}`) + out, err := tool.InvokableRun(ctx, `{"query":""}`) if err != nil { t.Fatalf("InvokableRun(empty query): %v", err) } var envelope githubEnvelope - if err := json.Unmarshal([]byte(out), &envelope); err != nil { + if err = json.Unmarshal([]byte(out), &envelope); err != nil { t.Fatalf("decode empty result: %v", err) } if len(envelope.Results) != 0 || envelope.Error != "" { @@ -181,19 +182,19 @@ func TestGitHub_BuildByNameUsesPythonNodeParams(t *testing.T) { if github.defaults.TopN != 17 { t.Errorf("defaults.TopN = %d, want 17", github.defaults.TopN) } - if _, err := BuildByName("github", map[string]any{"top_n": 100}); err != nil { + if _, err = BuildByName("github", map[string]any{"top_n": 100}); err != nil { t.Errorf("BuildByName(github) rejected GitHub's maximum top_n: %v", err) } - if _, err := BuildByName("github", map[string]any{"top_n": 0}); err == nil { + if _, err = BuildByName("github", map[string]any{"top_n": 0}); err == nil { t.Fatal("BuildByName(github) accepted non-positive top_n") } - if _, err := BuildByName("github", map[string]any{"top_n": 1.5}); err == nil { + if _, err = BuildByName("github", map[string]any{"top_n": 1.5}); err == nil { t.Fatal("BuildByName(github) accepted fractional top_n") } - if _, err := BuildByName("github", map[string]any{"top_n": "10"}); err == nil { + if _, err = BuildByName("github", map[string]any{"top_n": "10"}); err == nil { t.Fatal("BuildByName(github) accepted string top_n") } - if _, err := BuildByName("github", map[string]any{"top_n": 101}); err == nil { + if _, err = BuildByName("github", map[string]any{"top_n": 101}); err == nil { t.Fatal("BuildByName(github) accepted top_n above GitHub's per_page limit") } ignored, err := BuildByName("github", map[string]any{"max_results": 5}) @@ -220,6 +221,7 @@ func TestGitHub_ComponentContractMatchesPython(t *testing.T) { } func TestGitHub_ReferencesAndOutputsPreserveRawResults(t *testing.T) { + ctx := t.Context() github := NewGitHubTool() results := []any{map[string]any{ "name": "ragflow", @@ -234,7 +236,7 @@ func TestGitHub_ReferencesAndOutputsPreserveRawResults(t *testing.T) { } envelope := map[string]any{"results": results} - chunks, docAggs := github.BuildReferences(context.Background(), envelope) + chunks, docAggs := github.BuildReferences(ctx, envelope) if len(chunks) != 1 || len(docAggs) != 1 { t.Fatalf("references = %#v / %#v", chunks, docAggs) } @@ -269,9 +271,10 @@ func TestGitHub_LimitReferencesKeepsBoundaryChunk(t *testing.T) { func TestGitHub_Info(t *testing.T) { t.Parallel() + ctx := t.Context() tool := NewGitHubTool() - info, err := tool.Info(context.Background()) + info, err := tool.Info(ctx) if err != nil { t.Fatalf("Info: %v", err) } diff --git a/internal/storage/minio_test.go b/internal/storage/minio_test.go index 6eb4ee335f..5c69290d8f 100644 --- a/internal/storage/minio_test.go +++ b/internal/storage/minio_test.go @@ -20,7 +20,6 @@ package storage import ( "bytes" - "context" "fmt" "log" "ragflow/internal/utility" @@ -112,8 +111,9 @@ func TestNewMinioStorage_InvalidConfig(t *testing.T) { func TestMinioStorage_Health(t *testing.T) { storage := newTestMinioStorage(t) + ctx := t.Context() - healthy := storage.Health(context.Background()) + healthy := storage.Health(ctx) // Health check should return true if connection is working // Note: This depends on whether a default bucket is configured t.Logf("Health check result: %v", healthy) @@ -128,15 +128,16 @@ func TestMinioStorage_PutAndGet(t *testing.T) { bucket := "test-bucket" key := "test-file.txt" content := []byte("Hello, MinIO Test!") + ctx := t.Context() // Test Put - err := storage.Put(context.Background(), bucket, key, content) + err := storage.Put(ctx, bucket, key, content) if err != nil { t.Fatalf("Failed to put object: %v", err) } // Test Get - retrieved, err := storage.Get(context.Background(), bucket, key) + retrieved, err := storage.Get(ctx, bucket, key) if err != nil { t.Fatalf("Failed to get object: %v", err) } @@ -146,7 +147,7 @@ func TestMinioStorage_PutAndGet(t *testing.T) { } // Cleanup - err = storage.Remove(context.Background(), bucket, key) + err = storage.Remove(ctx, bucket, key) if err != nil { t.Logf("Warning: failed to cleanup test object: %v", err) } @@ -154,28 +155,30 @@ func TestMinioStorage_PutAndGet(t *testing.T) { func TestMinioStorage_Put_EmptyData(t *testing.T) { storage := newTestMinioStorage(t) + ctx := t.Context() bucket := "test-bucket" key := "empty-file.txt" content := []byte{} - err := storage.Put(context.Background(), bucket, key, content) + err := storage.Put(ctx, bucket, key, content) if err != nil { t.Fatalf("Failed to put empty object: %v", err) } // Verify object exists - exists := storage.ObjExist(context.Background(), bucket, key) + exists := storage.ObjExist(ctx, bucket, key) if !exists { t.Error("Expected empty object to exist") } // Cleanup - storage.Remove(context.Background(), bucket, key) + storage.Remove(ctx, bucket, key) } func TestMinioStorage_Put_LargeData(t *testing.T) { storage := newTestMinioStorage(t) + ctx := t.Context() bucket := "test-bucket" key := "large-file.bin" @@ -185,12 +188,12 @@ func TestMinioStorage_Put_LargeData(t *testing.T) { content[i] = byte(i % 256) } - err := storage.Put(context.Background(), bucket, key, content) + err := storage.Put(ctx, bucket, key, content) if err != nil { t.Fatalf("Failed to put large object: %v", err) } - retrieved, err := storage.Get(context.Background(), bucket, key) + retrieved, err := storage.Get(ctx, bucket, key) if err != nil { t.Fatalf("Failed to get large object: %v", err) } @@ -200,16 +203,17 @@ func TestMinioStorage_Put_LargeData(t *testing.T) { } // Cleanup - storage.Remove(context.Background(), bucket, key) + storage.Remove(ctx, bucket, key) } func TestMinioStorage_Get_NonExistent(t *testing.T) { storage := newTestMinioStorage(t) + ctx := t.Context() bucket := "test-bucket" key := "non-existent-file.txt" - _, err := storage.Get(context.Background(), bucket, key) + _, err := storage.Get(ctx, bucket, key) if err == nil { t.Error("Expected error when getting non-existent object") } @@ -217,31 +221,32 @@ func TestMinioStorage_Get_NonExistent(t *testing.T) { func TestMinioStorage_Remove(t *testing.T) { storage := newTestMinioStorage(t) + ctx := t.Context() bucket := "test-bucket" key := "file-to-delete.txt" content := []byte("Delete me") // First, put an object - err := storage.Put(context.Background(), bucket, key, content) + err := storage.Put(ctx, bucket, key, content) if err != nil { t.Fatalf("Failed to put object: %v", err) } // Verify it exists - exists := storage.ObjExist(context.Background(), bucket, key) + exists := storage.ObjExist(ctx, bucket, key) if !exists { t.Fatal("Expected object to exist before removal") } // Remove it - err = storage.Remove(context.Background(), bucket, key) + err = storage.Remove(ctx, bucket, key) if err != nil { t.Fatalf("Failed to remove object: %v", err) } // Verify it's gone - exists = storage.ObjExist(context.Background(), bucket, key) + exists = storage.ObjExist(ctx, bucket, key) if exists { t.Error("Expected object to not exist after removal") } @@ -249,12 +254,13 @@ func TestMinioStorage_Remove(t *testing.T) { func TestMinioStorage_Remove_NonExistent(t *testing.T) { storage := newTestMinioStorage(t) + ctx := t.Context() bucket := "test-bucket" key := "non-existent-file.txt" // Removing a non-existent object should not error - err := storage.Remove(context.Background(), bucket, key) + err := storage.Remove(ctx, bucket, key) if err != nil { t.Logf("Remove non-existent object returned error (may be acceptable): %v", err) } @@ -262,48 +268,50 @@ func TestMinioStorage_Remove_NonExistent(t *testing.T) { func TestMinioStorage_ObjExist(t *testing.T) { storage := newTestMinioStorage(t) + ctx := t.Context() bucket := "test-bucket" key := "existence-test.txt" content := []byte("Test content") // Check non-existent object - exists := storage.ObjExist(context.Background(), bucket, key) + exists := storage.ObjExist(ctx, bucket, key) if exists { t.Error("Expected non-existent object to return false") } // Create object - err := storage.Put(context.Background(), bucket, key, content) + err := storage.Put(ctx, bucket, key, content) if err != nil { t.Fatalf("Failed to put object: %v", err) } // Check existing object - exists = storage.ObjExist(context.Background(), bucket, key) + exists = storage.ObjExist(ctx, bucket, key) if !exists { t.Error("Expected existing object to return true") } // Cleanup - storage.Remove(context.Background(), bucket, key) + storage.Remove(ctx, bucket, key) } func TestMinioStorage_GetPresignedURL(t *testing.T) { storage := newTestMinioStorage(t) + ctx := t.Context() bucket := "test-bucket" key := "presigned-test.txt" content := []byte("Presigned URL test content") // Create object first - err := storage.Put(context.Background(), bucket, key, content) + err := storage.Put(ctx, bucket, key, content) if err != nil { t.Fatalf("Failed to put object: %v", err) } // Get presigned URL - url, err := storage.GetPresignedURL(context.Background(), bucket, key, 5*time.Minute) + url, err := storage.GetPresignedURL(ctx, bucket, key, 5*time.Minute) if err != nil { t.Fatalf("Failed to get presigned URL: %v", err) } @@ -318,16 +326,17 @@ func TestMinioStorage_GetPresignedURL(t *testing.T) { } // Cleanup - storage.Remove(context.Background(), bucket, key) + storage.Remove(ctx, bucket, key) } func TestMinioStorage_GetPresignedURL_NonExistent(t *testing.T) { storage := newTestMinioStorage(t) + ctx := t.Context() bucket := "test-bucket" key := "non-existent-presigned.txt" - _, err := storage.GetPresignedURL(context.Background(), bucket, key, 5*time.Minute) + _, err := storage.GetPresignedURL(ctx, bucket, key, 5*time.Minute) if err == nil { t.Log("Note: Some MinIO versions may allow presigned URLs for non-existent objects") } @@ -335,61 +344,63 @@ func TestMinioStorage_GetPresignedURL_NonExistent(t *testing.T) { func TestMinioStorage_BucketExists(t *testing.T) { storage := newTestMinioStorage(t) + ctx := t.Context() bucket := fmt.Sprintf("test-bucket-exists-%d", time.Now().Unix()) // Check non-existent bucket - exists := storage.BucketExists(context.Background(), bucket) + exists := storage.BucketExists(ctx, bucket) if exists { t.Error("Expected non-existent bucket to return false") } // Create bucket by putting an object - err := storage.Put(context.Background(), bucket, "test.txt", []byte("test")) + err := storage.Put(ctx, bucket, "test.txt", []byte("test")) if err != nil { t.Fatalf("Failed to create bucket: %v", err) } // Check existing bucket - exists = storage.BucketExists(context.Background(), bucket) + exists = storage.BucketExists(ctx, bucket) if !exists { t.Error("Expected existing bucket to return true") } // Cleanup - storage.RemoveBucket(context.Background(), bucket) + storage.RemoveBucket(ctx, bucket) } func TestMinioStorage_RemoveBucket(t *testing.T) { storage := newTestMinioStorage(t) + ctx := t.Context() bucket := fmt.Sprintf("test-bucket-remove-%d", time.Now().Unix()) // Create bucket with some objects - err := storage.Put(context.Background(), bucket, "file1.txt", []byte("content1")) + err := storage.Put(ctx, bucket, "file1.txt", []byte("content1")) if err != nil { t.Fatalf("Failed to put object: %v", err) } - err = storage.Put(context.Background(), bucket, "file2.txt", []byte("content2")) + err = storage.Put(ctx, bucket, "file2.txt", []byte("content2")) if err != nil { t.Fatalf("Failed to put object: %v", err) } // Verify bucket exists - exists := storage.BucketExists(context.Background(), bucket) + exists := storage.BucketExists(ctx, bucket) if !exists { t.Fatal("Expected bucket to exist before removal") } // Remove bucket - err = storage.RemoveBucket(context.Background(), bucket) + err = storage.RemoveBucket(ctx, bucket) if err != nil { t.Fatalf("Failed to remove bucket: %v", err) } // Verify bucket is gone - exists = storage.BucketExists(context.Background(), bucket) + exists = storage.BucketExists(ctx, bucket) if exists { t.Error("Expected bucket to not exist after removal") } @@ -397,6 +408,7 @@ func TestMinioStorage_RemoveBucket(t *testing.T) { func TestMinioStorage_Copy(t *testing.T) { storage := newTestMinioStorage(t) + ctx := t.Context() srcBucket := "test-bucket-src" srcKey := "source-file.txt" @@ -405,25 +417,25 @@ func TestMinioStorage_Copy(t *testing.T) { content := []byte("Content to copy") // Create source object - err := storage.Put(context.Background(), srcBucket, srcKey, content) + err := storage.Put(ctx, srcBucket, srcKey, content) if err != nil { t.Fatalf("Failed to put source object: %v", err) } // Copy object - success := storage.Copy(context.Background(), srcBucket, srcKey, destBucket, destKey) + success := storage.Copy(ctx, srcBucket, srcKey, destBucket, destKey) if !success { t.Fatal("Failed to copy object") } // Verify destination exists - exists := storage.ObjExist(context.Background(), destBucket, destKey) + exists := storage.ObjExist(ctx, destBucket, destKey) if !exists { t.Error("Expected copied object to exist") } // Verify content matches - retrieved, err := storage.Get(context.Background(), destBucket, destKey) + retrieved, err := storage.Get(ctx, destBucket, destKey) if err != nil { t.Fatalf("Failed to get copied object: %v", err) } @@ -433,33 +445,35 @@ func TestMinioStorage_Copy(t *testing.T) { } // Cleanup - storage.Remove(context.Background(), srcBucket, srcKey) - storage.Remove(context.Background(), destBucket, destKey) + storage.Remove(ctx, srcBucket, srcKey) + storage.Remove(ctx, destBucket, destKey) } func TestMinioStorage_Copy_NonExistentSource(t *testing.T) { storage := newTestMinioStorage(t) + ctx := t.Context() srcBucket := "test-bucket-src" srcKey := "non-existent-source.txt" destBucket := "test-bucket-dest" destKey := "should-not-exist.txt" - success := storage.Copy(context.Background(), srcBucket, srcKey, destBucket, destKey) + success := storage.Copy(ctx, srcBucket, srcKey, destBucket, destKey) if success { t.Error("Expected copy of non-existent object to fail") } // Verify destination does not exist - exists := storage.ObjExist(context.Background(), destBucket, destKey) + exists := storage.ObjExist(ctx, destBucket, destKey) if exists { t.Error("Expected destination object to not exist after failed copy") - storage.Remove(context.Background(), destBucket, destKey) + storage.Remove(ctx, destBucket, destKey) } } func TestMinioStorage_Move(t *testing.T) { storage := newTestMinioStorage(t) + ctx := t.Context() srcBucket := "test-bucket-src" srcKey := "file-to-move.txt" @@ -468,31 +482,31 @@ func TestMinioStorage_Move(t *testing.T) { content := []byte("Content to move") // Create source object - err := storage.Put(context.Background(), srcBucket, srcKey, content) + err := storage.Put(ctx, srcBucket, srcKey, content) if err != nil { t.Fatalf("Failed to put source object: %v", err) } // Move object - success := storage.Move(context.Background(), srcBucket, srcKey, destBucket, destKey) + success := storage.Move(ctx, srcBucket, srcKey, destBucket, destKey) if !success { t.Fatal("Failed to move object") } // Verify source is gone - exists := storage.ObjExist(context.Background(), srcBucket, srcKey) + exists := storage.ObjExist(ctx, srcBucket, srcKey) if exists { t.Error("Expected source object to not exist after move") } // Verify destination exists - exists = storage.ObjExist(context.Background(), destBucket, destKey) + exists = storage.ObjExist(ctx, destBucket, destKey) if !exists { t.Error("Expected moved object to exist") } // Verify content matches - retrieved, err := storage.Get(context.Background(), destBucket, destKey) + retrieved, err := storage.Get(ctx, destBucket, destKey) if err != nil { t.Fatalf("Failed to get moved object: %v", err) } @@ -502,18 +516,19 @@ func TestMinioStorage_Move(t *testing.T) { } // Cleanup - storage.Remove(context.Background(), destBucket, destKey) + storage.Remove(ctx, destBucket, destKey) } func TestMinioStorage_Move_NonExistentSource(t *testing.T) { storage := newTestMinioStorage(t) + ctx := t.Context() srcBucket := "test-bucket-src" srcKey := "non-existent-source.txt" destBucket := "test-bucket-dest" destKey := "should-not-exist.txt" - success := storage.Move(context.Background(), srcBucket, srcKey, destBucket, destKey) + success := storage.Move(ctx, srcBucket, srcKey, destBucket, destKey) if success { t.Error("Expected move of non-existent object to fail") } @@ -521,6 +536,7 @@ func TestMinioStorage_Move_NonExistentSource(t *testing.T) { func TestMinioStorage_MultipleObjectsInBucket(t *testing.T) { storage := newTestMinioStorage(t) + ctx := t.Context() bucket := fmt.Sprintf("test-multi-%d", time.Now().Unix()) numObjects := 10 @@ -529,7 +545,7 @@ func TestMinioStorage_MultipleObjectsInBucket(t *testing.T) { for i := 0; i < numObjects; i++ { key := fmt.Sprintf("file-%d.txt", i) content := []byte(fmt.Sprintf("Content %d", i)) - err := storage.Put(context.Background(), bucket, key, content) + err := storage.Put(ctx, bucket, key, content) if err != nil { t.Fatalf("Failed to put object %d: %v", i, err) } @@ -538,7 +554,7 @@ func TestMinioStorage_MultipleObjectsInBucket(t *testing.T) { // Verify all objects exist for i := 0; i < numObjects; i++ { key := fmt.Sprintf("file-%d.txt", i) - exists := storage.ObjExist(context.Background(), bucket, key) + exists := storage.ObjExist(ctx, bucket, key) if !exists { t.Errorf("Expected object %s to exist", key) } @@ -548,7 +564,7 @@ func TestMinioStorage_MultipleObjectsInBucket(t *testing.T) { for i := 0; i < numObjects; i++ { key := fmt.Sprintf("file-%d.txt", i) expectedContent := []byte(fmt.Sprintf("Content %d", i)) - retrieved, err := storage.Get(context.Background(), bucket, key) + retrieved, err := storage.Get(ctx, bucket, key) if err != nil { t.Errorf("Failed to get object %s: %v", key, err) continue @@ -559,7 +575,7 @@ func TestMinioStorage_MultipleObjectsInBucket(t *testing.T) { } // Cleanup - remove bucket with all objects - err := storage.RemoveBucket(context.Background(), bucket) + err := storage.RemoveBucket(ctx, bucket) if err != nil { t.Logf("Warning: failed to cleanup bucket: %v", err) } @@ -567,6 +583,7 @@ func TestMinioStorage_MultipleObjectsInBucket(t *testing.T) { func TestMinioStorage_SpecialCharactersInKey(t *testing.T) { storage := newTestMinioStorage(t) + ctx := t.Context() bucket := "test-bucket" specialKeys := []string{ @@ -581,13 +598,13 @@ func TestMinioStorage_SpecialCharactersInKey(t *testing.T) { for _, key := range specialKeys { content := []byte(fmt.Sprintf("Content for %s", key)) - err := storage.Put(context.Background(), bucket, key, content) + err := storage.Put(ctx, bucket, key, content) if err != nil { t.Errorf("Failed to put object with key '%s': %v", key, err) continue } - retrieved, err := storage.Get(context.Background(), bucket, key) + retrieved, err := storage.Get(ctx, bucket, key) if err != nil { t.Errorf("Failed to get object with key '%s': %v", key, err) continue @@ -598,12 +615,13 @@ func TestMinioStorage_SpecialCharactersInKey(t *testing.T) { } // Cleanup - storage.Remove(context.Background(), bucket, key) + storage.Remove(ctx, bucket, key) } } func TestMinioStorage_TenantID(t *testing.T) { storage := newTestMinioStorage(t) + ctx := t.Context() bucket := "test-bucket" key := "tenant-test.txt" @@ -611,13 +629,13 @@ func TestMinioStorage_TenantID(t *testing.T) { tenantID := "tenant-123" // Put with tenant ID - err := storage.Put(context.Background(), bucket, key, content, tenantID) + err := storage.Put(ctx, bucket, key, content, tenantID) if err != nil { t.Fatalf("Failed to put object with tenant ID: %v", err) } // Get with tenant ID - retrieved, err := storage.Get(context.Background(), bucket, key, tenantID) + retrieved, err := storage.Get(ctx, bucket, key, tenantID) if err != nil { t.Fatalf("Failed to get object with tenant ID: %v", err) } @@ -627,11 +645,11 @@ func TestMinioStorage_TenantID(t *testing.T) { } // Check existence with tenant ID - exists := storage.ObjExist(context.Background(), bucket, key, tenantID) + exists := storage.ObjExist(ctx, bucket, key, tenantID) if !exists { t.Error("Expected object to exist with tenant ID") } // Cleanup - storage.Remove(context.Background(), bucket, key, tenantID) + storage.Remove(ctx, bucket, key, tenantID) }