From 4386ff71b02ef3adb43335590c089bf2ea64ac96 Mon Sep 17 00:00:00 2001 From: Jin Hai Date: Tue, 11 Aug 2026 20:12:29 +0800 Subject: [PATCH] Go: fix context, part2 (#18133) Signed-off-by: Jin Hai --- internal/agent/tool/google_scholar_test.go | 18 ++++--- internal/agent/tool/google_test.go | 13 +++-- internal/agent/tool/http_helper_test.go | 50 +++++++++++-------- internal/agent/tool/jin10_test.go | 10 ++-- internal/agent/tool/keenable_test.go | 37 +++++++++----- internal/agent/tool/mcp_test.go | 9 ++-- internal/agent/tool/pubmed_test.go | 16 ++++--- internal/agent/tool/qweather_test.go | 16 ++++--- internal/agent/tool/retrieval_test.go | 38 +++++++++------ internal/agent/tool/searxng_test.go | 21 ++++---- internal/agent/tool/tavily_test.go | 56 ++++++++++++++-------- internal/agent/tool/tushare_test.go | 16 ++++--- internal/agent/tool/wencai_test.go | 12 +++-- internal/agent/tool/wikipedia_test.go | 15 +++--- internal/agent/tool/yahoo_finance_test.go | 38 +++++++++------ 15 files changed, 227 insertions(+), 138 deletions(-) diff --git a/internal/agent/tool/google_scholar_test.go b/internal/agent/tool/google_scholar_test.go index a13215c531..332757f616 100644 --- a/internal/agent/tool/google_scholar_test.go +++ b/internal/agent/tool/google_scholar_test.go @@ -17,7 +17,6 @@ package tool import ( - "context" "encoding/json" "net/http" "net/http/httptest" @@ -180,6 +179,7 @@ func TestGoogleScholar_BuildURL(t *testing.T) { func TestGoogleScholar_ParseResults(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") @@ -191,7 +191,7 @@ func TestGoogleScholar_ParseResults(t *testing.T) { Transport: rewriteHostTransport(srv.URL), }) tool := NewGoogleScholarToolWith(helper) - out, err := tool.InvokableRun(context.Background(), + out, err := tool.InvokableRun(ctx, `{"query":"transformer","top_n":5,"sort_by":"relevance","year_low":2020,"year_high":2024,"patents":true}`) if err != nil { t.Fatalf("InvokableRun: %v", err) @@ -229,23 +229,25 @@ func TestGoogleScholar_ParseResults(t *testing.T) { func TestGoogleScholar_EmptyQuery(t *testing.T) { t.Parallel() + ctx := t.Context() tool := NewGoogleScholarTool() - out, err := tool.InvokableRun(context.Background(), `{"query":""}`) + out, err := tool.InvokableRun(ctx, `{"query":""}`) if err != nil { t.Fatalf("InvokableRun(empty): %v", err) } var envelope googleScholarEnvelope - 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 TestGoogleScholar_Info(t *testing.T) { t.Parallel() + ctx := t.Context() tool := NewGoogleScholarTool() - info, err := tool.Info(context.Background()) + info, err := tool.Info(ctx) if err != nil { t.Fatalf("Info: %v", err) } @@ -259,6 +261,7 @@ func TestGoogleScholar_Info(t *testing.T) { func TestGoogleScholar_ComponentReferencesAndValidation(t *testing.T) { t.Parallel() + ctx := t.Context() built, err := BuildByName("google_scholar", map[string]any{ "top_n": float64(7), @@ -286,7 +289,7 @@ func TestGoogleScholar_ComponentReferencesAndValidation(t *testing.T) { envelope := map[string]any{"results": []any{map[string]any{ "title": "Paper", "link": "https://paper.example", "authors": "A Author", "year": "2024", "snippet": "Abstract", }}} - chunks, docAggs := scholar.BuildReferences(context.Background(), envelope) + chunks, docAggs := scholar.BuildReferences(ctx, envelope) if len(chunks) != 1 || len(docAggs) != 1 || !strings.Contains(chunks[0]["content"].(string), "Authors: A Author") { t.Fatalf("references = %#v / %#v", chunks, docAggs) } @@ -325,6 +328,7 @@ func TestRenderGoogleScholarReferencesStopsBeforeOverBudgetBlock(t *testing.T) { func TestGoogleScholar_MergesNodeLevelDefaults(t *testing.T) { t.Parallel() + ctx := t.Context() srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { q := r.URL.Query() @@ -362,7 +366,7 @@ func TestGoogleScholar_MergesNodeLevelDefaults(t *testing.T) { YearHigh: 2024, Patents: boolPtr(false), }) - _, err := tool.InvokableRun(context.Background(), `{}`) + _, err := tool.InvokableRun(ctx, `{}`) if err != nil { t.Fatalf("InvokableRun with defaults: %v", err) } diff --git a/internal/agent/tool/google_test.go b/internal/agent/tool/google_test.go index 242a8323c0..232779539e 100644 --- a/internal/agent/tool/google_test.go +++ b/internal/agent/tool/google_test.go @@ -17,7 +17,6 @@ package tool import ( - "context" "encoding/json" "net/http" "net/http/httptest" @@ -86,6 +85,7 @@ func TestGoogle_BuildURL(t *testing.T) { func TestGoogle_ParseResults(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", "application/json") @@ -102,7 +102,7 @@ func TestGoogle_ParseResults(t *testing.T) { Transport: rewriteHostTransport(srv.URL), }) tool := NewGoogleToolWith(helper) - out, err := tool.InvokableRun(context.Background(), + out, err := tool.InvokableRun(ctx, `{"q":"ragflow","api_key":"K","num":5,"country":"us","language":"en"}`) if err != nil { t.Fatalf("InvokableRun: %v", err) @@ -128,9 +128,10 @@ func TestGoogle_ParseResults(t *testing.T) { func TestGoogle_RequiresAPIKey(t *testing.T) { t.Parallel() + ctx := t.Context() tool := NewGoogleTool() - _, err := tool.InvokableRun(context.Background(), `{"q":"x","api_key":""}`) + _, err := tool.InvokableRun(ctx, `{"q":"x","api_key":""}`) if err == nil { t.Fatal("expected error for missing api_key") } @@ -141,9 +142,10 @@ func TestGoogle_RequiresAPIKey(t *testing.T) { func TestGoogle_InfoAndInputForm(t *testing.T) { t.Parallel() + ctx := t.Context() tool := NewGoogleTool() - info, err := tool.Info(context.Background()) + info, err := tool.Info(ctx) if err != nil { t.Fatalf("Info: %v", err) } @@ -246,6 +248,7 @@ func TestGoogle_BuildByNameIgnoresUnrelatedCanvasParams(t *testing.T) { func TestGoogle_ComponentReferencesAndOutputs(t *testing.T) { t.Parallel() + ctx := t.Context() google := NewGoogleTool() spec := google.ComponentSpec() @@ -270,7 +273,7 @@ func TestGoogle_ComponentReferencesAndOutputs(t *testing.T) { }}, }, }} - chunks, docAggs := google.BuildReferences(context.Background(), envelope) + chunks, docAggs := google.BuildReferences(ctx, envelope) if len(chunks) != 2 || len(docAggs) != 2 { t.Fatalf("references = %#v / %#v", chunks, docAggs) } diff --git a/internal/agent/tool/http_helper_test.go b/internal/agent/tool/http_helper_test.go index 8923628ced..369289ab7b 100644 --- a/internal/agent/tool/http_helper_test.go +++ b/internal/agent/tool/http_helper_test.go @@ -49,6 +49,7 @@ func newTestHelper(maxAttempts int, base, max time.Duration) *HTTPHelper { // attempt with no retry, and the body / content-type round-trip cleanly. func TestHTTPHelper_HappyPath(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", "application/json") @@ -58,7 +59,7 @@ func TestHTTPHelper_HappyPath(t *testing.T) { defer srv.Close() h := newTestHelper(3, 1*time.Millisecond, 5*time.Millisecond) - resp, err := h.Do(context.Background(), http.MethodGet, srv.URL, "", "", nil) + resp, err := h.Do(ctx, http.MethodGet, srv.URL, "", "", nil) if err != nil { t.Fatalf("Do returned error: %v", err) } @@ -80,6 +81,7 @@ func TestHTTPHelper_HappyPath(t *testing.T) { // returns the first 2xx response. Server returns 503 twice, then 200. func TestHTTPHelper_RetriesOn5xx(t *testing.T) { t.Parallel() + ctx := t.Context() var hits int32 srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { @@ -94,7 +96,7 @@ func TestHTTPHelper_RetriesOn5xx(t *testing.T) { defer srv.Close() h := newTestHelper(3, 1*time.Millisecond, 5*time.Millisecond) - resp, err := h.Do(context.Background(), http.MethodGet, srv.URL, "", "", nil) + resp, err := h.Do(ctx, http.MethodGet, srv.URL, "", "", nil) if err != nil { t.Fatalf("Do returned error: %v", err) } @@ -116,6 +118,7 @@ func TestHTTPHelper_RetriesOn5xx(t *testing.T) { // no retry — the caller is responsible for fixing 4xx, retrying won't help. func TestHTTPHelper_NoRetryOn4xx(t *testing.T) { t.Parallel() + ctx := t.Context() var hits int32 srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { @@ -126,7 +129,7 @@ func TestHTTPHelper_NoRetryOn4xx(t *testing.T) { defer srv.Close() h := newTestHelper(3, 1*time.Millisecond, 5*time.Millisecond) - resp, err := h.Do(context.Background(), http.MethodGet, srv.URL, "", "", nil) + resp, err := h.Do(ctx, http.MethodGet, srv.URL, "", "", nil) if err != nil { t.Fatalf("Do returned error: %v", err) } @@ -158,6 +161,7 @@ func TestHTTPHelper_NoRetryOn4xx(t *testing.T) { // and flaked under load. func TestHTTPHelper_Timeout(t *testing.T) { t.Parallel() + ctx := t.Context() var hits int32 srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { @@ -176,7 +180,7 @@ func TestHTTPHelper_Timeout(t *testing.T) { }) // Tight 50ms deadline. The server takes 500ms, so this call must // abort due to the context, not the server finishing. - ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond) + ctx, cancel := context.WithTimeout(ctx, 50*time.Millisecond) defer cancel() _, err := h.Do(ctx, http.MethodGet, srv.URL, "", "", nil) @@ -196,6 +200,7 @@ func TestHTTPHelper_Timeout(t *testing.T) { // 5xx error is returned. func TestHTTPHelper_5xxExhaustion(t *testing.T) { t.Parallel() + ctx := t.Context() var hits int32 srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { @@ -205,7 +210,7 @@ func TestHTTPHelper_5xxExhaustion(t *testing.T) { defer srv.Close() h := newTestHelper(3, 1*time.Millisecond, 5*time.Millisecond) - _, err := h.Do(context.Background(), http.MethodGet, srv.URL, "", "", nil) + _, err := h.Do(ctx, http.MethodGet, srv.URL, "", "", nil) if err == nil { t.Fatal("expected error after 5xx exhaustion, got nil") } @@ -218,6 +223,7 @@ func TestHTTPHelper_5xxExhaustion(t *testing.T) { // custom headers and a non-empty content-type on POST bodies. func TestHTTPHelper_HeadersAndContentType(t *testing.T) { t.Parallel() + ctx := t.Context() srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if got := r.Header.Get("X-Token"); got != "abc" { @@ -234,7 +240,7 @@ func TestHTTPHelper_HeadersAndContentType(t *testing.T) { defer srv.Close() h := newTestHelper(1, 1*time.Millisecond, 5*time.Millisecond) - resp, err := h.Do(context.Background(), http.MethodPost, srv.URL, `{"k":1}`, "application/json", map[string]string{"X-Token": "abc"}) + resp, err := h.Do(ctx, http.MethodPost, srv.URL, `{"k":1}`, "application/json", map[string]string{"X-Token": "abc"}) if err != nil { t.Fatalf("Do: %v", err) } @@ -251,20 +257,20 @@ func TestBackoffExponential(t *testing.T) { t.Parallel() base := 50 * time.Millisecond - max := 300 * time.Millisecond + maxDuration := 300 * time.Millisecond - got1 := backoff(base, max, 1) + got1 := backoff(base, maxDuration, 1) if got1 < 0 || got1 > base { t.Fatalf("backoff(attempt=1) = %s, want [0, %s]", got1, base) } - got3 := backoff(base, max, 3) - if got3 < 0 || got3 > max { - t.Fatalf("backoff(attempt=3) = %s, want [0, %s] (capped)", got3, max) + got3 := backoff(base, maxDuration, 3) + if got3 < 0 || got3 > maxDuration { + t.Fatalf("backoff(attempt=3) = %s, want [0, %s] (capped)", got3, maxDuration) } // With base=50ms, attempt=10 should be capped at 300ms. - got10 := backoff(base, max, 10) - if got10 > max { - t.Fatalf("backoff(attempt=10) = %s, want <= %s (cap)", got10, max) + got10 := backoff(base, maxDuration, 10) + if got10 > maxDuration { + t.Fatalf("backoff(attempt=10) = %s, want <= %s (cap)", got10, maxDuration) } } @@ -289,7 +295,7 @@ func TestRetryConfigDefaults(t *testing.T) { // test for the M1-rebinding fix as hardened by the post-Phase-7 // review: DNS pinning MUST happen at the transport layer, not by // rewriting the request URL. If the URL host were rewritten to the IP, -// the TLS ServerName (auto-populated by Go from req.URL.Host) would +// the TLS ServerName (autopopulated by Go from req.URL.Host) would // become the IP, the SNI would send the IP, and cert verification // would target the IP — which is not what real HTTPS sites have, and // would manifest as x509 errors against any host-cert-only target. @@ -297,12 +303,13 @@ func TestRetryConfigDefaults(t *testing.T) { // This test stands up a real TLS server with a cert whose DNS SAN is // "example.test" and whose IP SAN covers the loopback address. The // pinned dialer connects to 127.0.0.1, but the request URL host stays -// as "example.test". The server observes the SNI the client sent and +// as "example.test". The server observes the SNI the client sent, and // we assert it equals "example.test" (not the IP), and the request // completes successfully (cert verification passes because the URL // host matches the SAN). func TestHTTPHelper_DoPinnedHTTPS_PreservesSNIAndCert(t *testing.T) { t.Parallel() + ctx := t.Context() // Cert valid for "example.test" (DNS SAN) and 127.0.0.1, ::1 // (IP SANs, just so the test environment itself can resolve). @@ -353,7 +360,7 @@ func TestHTTPHelper_DoPinnedHTTPS_PreservesSNIAndCert(t *testing.T) { h := NewHTTPHelper() h.baseTransport.TLSClientConfig = &tls.Config{RootCAs: pool} - resp, err := h.DoPinned(context.Background(), + resp, err := h.DoPinned(ctx, http.MethodGet, targetURL, "", "", nil, "example.test", pinnedIP) if err != nil { t.Fatalf("DoPinned: %v", err) @@ -378,9 +385,10 @@ func TestHTTPHelper_DoPinnedHTTPS_PreservesSNIAndCert(t *testing.T) { // refuse rather than silently deliver a broken connection. func TestHTTPHelper_DoPinnedRefusesMismatchedURLHost(t *testing.T) { t.Parallel() + ctx := t.Context() h := NewHTTPHelper() - _, err := h.DoPinned(context.Background(), + _, err := h.DoPinned(ctx, http.MethodGet, "https://attacker.example/foo", "", "", nil, @@ -414,6 +422,7 @@ func TestHTTPHelper_DoPinnedRefusesMismatchedURLHost(t *testing.T) { // bypasses the proxy, the direct dial to 127.0.0.1 succeeds. func TestHTTPHelper_DoPinnedBypassesProxy(t *testing.T) { t.Parallel() + ctx := t.Context() cert := generateTestCert(t, "example.test") srv := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { @@ -467,7 +476,7 @@ func TestHTTPHelper_DoPinnedBypassesProxy(t *testing.T) { } h.baseTransport.Proxy = http.ProxyURL(closedProxy) - resp, err := h.DoPinned(context.Background(), + resp, err := h.DoPinned(ctx, http.MethodGet, targetURL, "", "", nil, "example.test", pinnedIP) if err != nil { t.Fatalf("DoPinned: %v (proxy may not have been bypassed)", err) @@ -486,6 +495,7 @@ func TestHTTPHelper_DoPinnedBypassesProxy(t *testing.T) { // transport-layer primitive that DoPinned uses. func TestPinnedDialer_RewritesAddress(t *testing.T) { t.Parallel() + ctx := t.Context() // Stand up a TCP server on 127.0.0.1: and capture the // accepted conn to confirm the pinned dialer actually dialed it. @@ -511,7 +521,7 @@ func TestPinnedDialer_RewritesAddress(t *testing.T) { // Pass a deliberately misleading host in the addr — the dialer // must ignore it and dial 127.0.0.1: instead. misleading := net.JoinHostPort("203.0.113.99", fmt.Sprint(ln.Addr().(*net.TCPAddr).Port)) - conn, derr := d.DialContext(context.Background(), "tcp", misleading) + conn, derr := d.DialContext(ctx, "tcp", misleading) if derr != nil { t.Fatalf("DialContext: %v", derr) } diff --git a/internal/agent/tool/jin10_test.go b/internal/agent/tool/jin10_test.go index 26b86312e2..d88b208612 100644 --- a/internal/agent/tool/jin10_test.go +++ b/internal/agent/tool/jin10_test.go @@ -17,7 +17,6 @@ package tool import ( - "context" "encoding/json" "strings" "testing" @@ -25,6 +24,7 @@ import ( func TestJin10_StubsUnsupported(t *testing.T) { t.Parallel() + ctx := t.Context() cases := []struct { name string @@ -47,7 +47,7 @@ func TestJin10_StubsUnsupported(t *testing.T) { t.Run(tc.name, func(t *testing.T) { t.Parallel() tool := NewJin10Tool() - out, err := tool.InvokableRun(context.Background(), tc.args) + out, err := tool.InvokableRun(ctx, tc.args) if err == nil { t.Fatalf("expected error, got nil (out=%s)", out) } @@ -70,9 +70,10 @@ func TestJin10_StubsUnsupported(t *testing.T) { func TestJin10_RejectsMalformedJSON(t *testing.T) { t.Parallel() + ctx := t.Context() tool := NewJin10Tool() - _, 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") } @@ -83,9 +84,10 @@ func TestJin10_RejectsMalformedJSON(t *testing.T) { func TestJin10_Info(t *testing.T) { t.Parallel() + ctx := t.Context() tool := NewJin10Tool() - 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/keenable_test.go b/internal/agent/tool/keenable_test.go index ca8d21e0c2..f5ccb49a1e 100644 --- a/internal/agent/tool/keenable_test.go +++ b/internal/agent/tool/keenable_test.go @@ -17,7 +17,6 @@ package tool import ( - "context" "encoding/json" "net/http" "net/http/httptest" @@ -30,6 +29,7 @@ import ( // without an X-API-Key header. func TestKeenable_KeylessPath(t *testing.T) { t.Parallel() + ctx := t.Context() var gotMethod, gotPath, gotUA, gotTitle, gotAPIKey, gotCT string var gotBody map[string]any @@ -52,7 +52,7 @@ func TestKeenable_KeylessPath(t *testing.T) { }) tool := NewKeenableToolWithEnvBaseURL(helper, func() string { return "https://" + srv.URL[len("http://"):] }) - if _, err := tool.InvokableRun(context.Background(), `{"query":"ragflow"}`); err != nil { + if _, err := tool.InvokableRun(ctx, `{"query":"ragflow"}`); err != nil { t.Fatalf("InvokableRun: %v", err) } @@ -87,6 +87,7 @@ func TestKeenable_KeylessPath(t *testing.T) { // the request. func TestKeenable_KeyedPath(t *testing.T) { t.Parallel() + ctx := t.Context() var gotPath, gotAPIKey string srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { @@ -103,7 +104,7 @@ func TestKeenable_KeyedPath(t *testing.T) { tool := NewKeenableToolWithAPIKey(helper, "key-xyz") tool.envBaseURL = func() string { return "https://" + srv.URL[len("http://"):] } - if _, err := tool.InvokableRun(context.Background(), + if _, err := tool.InvokableRun(ctx, `{"query":"ragflow","mode":"realtime"}`); err != nil { t.Fatalf("InvokableRun: %v", err) } @@ -120,6 +121,7 @@ func TestKeenable_KeyedPath(t *testing.T) { // that the result list is truncated to top_n. func TestKeenable_SiteAndTopN(t *testing.T) { t.Parallel() + ctx := t.Context() var gotBody map[string]any srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { @@ -139,7 +141,7 @@ func TestKeenable_SiteAndTopN(t *testing.T) { }) tool := NewKeenableToolWithEnvBaseURL(helper, func() string { return "https://" + srv.URL[len("http://"):] }) - out, err := tool.InvokableRun(context.Background(), + out, err := tool.InvokableRun(ctx, `{"query":"x","site":"example.com","top_n":2}`) if err != nil { t.Fatalf("InvokableRun: %v", err) @@ -171,6 +173,7 @@ func TestKeenable_SiteAndTopN(t *testing.T) { // results from the upstream response (the default in the Python tool). func TestKeenable_DefaultTopN(t *testing.T) { t.Parallel() + ctx := t.Context() srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { // 12 results; default top_n is 10, so we expect 10 in the envelope. @@ -193,7 +196,7 @@ func TestKeenable_DefaultTopN(t *testing.T) { }) tool := NewKeenableToolWithEnvBaseURL(helper, func() string { return "https://" + srv.URL[len("http://"):] }) - out, err := tool.InvokableRun(context.Background(), `{"query":"x"}`) + out, err := tool.InvokableRun(ctx, `{"query":"x"}`) if err != nil { t.Fatalf("InvokableRun: %v", err) } @@ -209,9 +212,10 @@ func TestKeenable_DefaultTopN(t *testing.T) { func TestKeenable_MissingQuery(t *testing.T) { t.Parallel() + ctx := t.Context() tool := NewKeenableTool() - out, err := tool.InvokableRun(context.Background(), `{}`) + out, err := tool.InvokableRun(ctx, `{}`) if err != nil { t.Fatalf("InvokableRun: %v", err) } @@ -228,9 +232,10 @@ func TestKeenable_MissingQuery(t *testing.T) { // of realtime mode without a configured api_key. func TestKeenable_RealtimeRequiresAPIKey(t *testing.T) { t.Parallel() + ctx := t.Context() tool := NewKeenableTool() - _, err := tool.InvokableRun(context.Background(), `{"query":"x","mode":"realtime"}`) + _, err := tool.InvokableRun(ctx, `{"query":"x","mode":"realtime"}`) if err == nil { t.Fatal("expected error for realtime mode without api_key") } @@ -243,9 +248,10 @@ func TestKeenable_RealtimeRequiresAPIKey(t *testing.T) { // up front instead of being forwarded to the upstream. func TestKeenable_InvalidMode(t *testing.T) { t.Parallel() + ctx := t.Context() tool := NewKeenableTool() - _, err := tool.InvokableRun(context.Background(), `{"query":"x","mode":"bogus"}`) + _, err := tool.InvokableRun(ctx, `{"query":"x","mode":"bogus"}`) if err == nil { t.Fatal("expected error for invalid mode") } @@ -302,6 +308,7 @@ func TestKeenable_ResolveBaseURL(t *testing.T) { // the test does not depend on the host environment. func TestKeenable_BaseURLFromEnv(t *testing.T) { t.Parallel() + ctx := t.Context() var gotPath string srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { @@ -318,7 +325,7 @@ func TestKeenable_BaseURLFromEnv(t *testing.T) { return "https://" + srv.URL[len("http://"):] }) - if _, err := tool.InvokableRun(context.Background(), `{"query":"x"}`); err != nil { + if _, err := tool.InvokableRun(ctx, `{"query":"x"}`); err != nil { t.Fatalf("InvokableRun: %v", err) } if gotPath != "/v1/search/public" { @@ -330,9 +337,10 @@ func TestKeenable_BaseURLFromEnv(t *testing.T) { // reported back to the caller instead of being silently sent. func TestKeenable_BadBaseURL(t *testing.T) { t.Parallel() + ctx := t.Context() tool := NewKeenableToolWithEnvBaseURL(NewHTTPHelper(), func() string { return "http://example.com" }) - _, err := tool.InvokableRun(context.Background(), `{"query":"x"}`) + _, err := tool.InvokableRun(ctx, `{"query":"x"}`) if err == nil { t.Fatal("expected error for non-https non-loopback base URL") } @@ -345,6 +353,7 @@ func TestKeenable_BadBaseURL(t *testing.T) { // is surfaced as an error and an _ERROR envelope. func TestKeenable_UpstreamError(t *testing.T) { t.Parallel() + ctx := t.Context() srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { http.Error(w, "boom", http.StatusInternalServerError) @@ -356,7 +365,7 @@ func TestKeenable_UpstreamError(t *testing.T) { }) tool := NewKeenableToolWithEnvBaseURL(helper, func() string { return "https://" + srv.URL[len("http://"):] }) - out, err := tool.InvokableRun(context.Background(), `{"query":"x"}`) + out, err := tool.InvokableRun(ctx, `{"query":"x"}`) if err == nil { t.Fatal("expected error for 5xx response") } @@ -372,9 +381,10 @@ func TestKeenable_UpstreamError(t *testing.T) { // TestKeenable_Info verifies the model-facing metadata. func TestKeenable_Info(t *testing.T) { t.Parallel() + ctx := t.Context() tool := NewKeenableTool() - info, err := tool.Info(context.Background()) + info, err := tool.Info(ctx) if err != nil { t.Fatalf("Info: %v", err) } @@ -411,6 +421,7 @@ func TestKeenable_Info(t *testing.T) { func TestKeenable_ComponentContractReferencesAndOutputs(t *testing.T) { t.Parallel() + ctx := t.Context() tool := NewKeenableTool() spec := tool.ComponentSpec() @@ -430,7 +441,7 @@ func TestKeenable_ComponentContractReferencesAndOutputs(t *testing.T) { "url": "https://example.com/item", "description": "Fresh search result", }}} - chunks, docAggs := tool.BuildReferences(context.Background(), envelope) + chunks, docAggs := tool.BuildReferences(ctx, envelope) if len(chunks) != 1 || len(docAggs) != 1 { t.Fatalf("references = %#v / %#v", chunks, docAggs) } diff --git a/internal/agent/tool/mcp_test.go b/internal/agent/tool/mcp_test.go index f7d1c4b1e9..149295a4a9 100644 --- a/internal/agent/tool/mcp_test.go +++ b/internal/agent/tool/mcp_test.go @@ -63,8 +63,9 @@ func TestMCPToolAdapter_InfoReturnsMCPDescriptor(t *testing.T) { // mcpclient is discovery-only; InvokableRun must return a clear error // until tools/call lands. func TestMCPToolAdapter_InvokableRunNotYetImplemented(t *testing.T) { + ctx := t.Context() a := NewMCPToolAdapter(mcpclient.Tool{Name: "x"}) - out, err := a.InvokableRun(context.Background(), `{"q":"hi"}`) + out, err := a.InvokableRun(ctx, `{"q":"hi"}`) if err == nil { t.Fatal("expected error from unimplemented tools/call") } @@ -162,6 +163,7 @@ func TestMarshalArguments_ValidJSON(t *testing.T) { // envelope (string result) and the session lifecycle // (initialize → tools/call). func TestMCPToolAdapter_InvokableRunDispatchesCallTool(t *testing.T) { + ctx := t.Context() defer mcpLoopbackOverride(t)() var sawCall bool srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { @@ -185,7 +187,7 @@ func TestMCPToolAdapter_InvokableRunDispatchesCallTool(t *testing.T) { defer srv.Close() a := NewMCPToolAdapterFull(mcpclient.Tool{Name: "echo"}, srv.URL, nil, 2*time.Second, srv.Client()) - out, err := a.InvokableRun(context.Background(), `{"msg":"hi"}`) + out, err := a.InvokableRun(ctx, `{"msg":"hi"}`) if err != nil { t.Fatalf("InvokableRun: %v", err) } @@ -200,6 +202,7 @@ func TestMCPToolAdapter_InvokableRunDispatchesCallTool(t *testing.T) { // TestMCPToolAdapter_InvokableRunIsError: a tools/call response // with isError=true surfaces as a Go error. func TestMCPToolAdapter_InvokableRunIsError(t *testing.T) { + ctx := t.Context() defer mcpLoopbackOverride(t)() srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { body, _ := io.ReadAll(r.Body) @@ -220,7 +223,7 @@ func TestMCPToolAdapter_InvokableRunIsError(t *testing.T) { defer srv.Close() a := NewMCPToolAdapterFull(mcpclient.Tool{Name: "echo"}, srv.URL, nil, 2*time.Second, srv.Client()) - _, err := a.InvokableRun(context.Background(), `{}`) + _, err := a.InvokableRun(ctx, `{}`) if err == nil { t.Fatalf("expected error for isError response") } diff --git a/internal/agent/tool/pubmed_test.go b/internal/agent/tool/pubmed_test.go index 071116c017..b218c454e0 100644 --- a/internal/agent/tool/pubmed_test.go +++ b/internal/agent/tool/pubmed_test.go @@ -17,7 +17,6 @@ package tool import ( - "context" "encoding/json" "net/http" "net/http/httptest" @@ -76,6 +75,7 @@ func TestPubMed_BuildURL(t *testing.T) { } func TestPubMed_InvokableRunParsesXML(t *testing.T) { + ctx := t.Context() t.Parallel() var esearchHits, efetchHits int32 @@ -123,7 +123,7 @@ func TestPubMed_InvokableRunParsesXML(t *testing.T) { helper := NewHTTPHelper().WithClient(&http.Client{Transport: rewriteHostTransport(srv.URL)}) tool := NewPubMedToolWithDefaults(helper, pubmedParams{TopN: 3, Email: "tester@example.com"}) - out, err := tool.InvokableRun(context.Background(), `{"query":"ragflow"}`) + out, err := tool.InvokableRun(ctx, `{"query":"ragflow"}`) if err != nil { t.Fatalf("InvokableRun: %v", err) } @@ -168,6 +168,7 @@ func TestPubMed_InvokableRunParsesXML(t *testing.T) { } func TestPubMed_InvokableRunEmptyResults(t *testing.T) { + ctx := t.Context() t.Parallel() srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { @@ -178,7 +179,7 @@ func TestPubMed_InvokableRunEmptyResults(t *testing.T) { helper := NewHTTPHelper().WithClient(&http.Client{Transport: rewriteHostTransport(srv.URL)}) tool := NewPubMedToolWith(helper) - out, err := tool.InvokableRun(context.Background(), `{"query":"missing"}`) + out, err := tool.InvokableRun(ctx, `{"query":"missing"}`) if err != nil { t.Fatalf("InvokableRun: %v", err) } @@ -192,10 +193,11 @@ func TestPubMed_InvokableRunEmptyResults(t *testing.T) { } func TestPubMed_InvokableRunEmptyQuery(t *testing.T) { + ctx := t.Context() t.Parallel() tool := NewPubMedTool() - out, err := tool.InvokableRun(context.Background(), `{"query":""}`) + out, err := tool.InvokableRun(ctx, `{"query":""}`) if err != nil { t.Fatalf("InvokableRun(empty): %v", err) } @@ -206,10 +208,11 @@ func TestPubMed_InvokableRunEmptyQuery(t *testing.T) { } func TestPubMed_InfoOnlyExposesQuery(t *testing.T) { + ctx := t.Context() t.Parallel() tool := NewPubMedTool() - info, err := tool.Info(context.Background()) + info, err := tool.Info(ctx) if err != nil { t.Fatalf("Info: %v", err) } @@ -258,6 +261,7 @@ func TestPubMed_BuildByNameAcceptsNodeParams(t *testing.T) { } func TestPubMed_ComponentReferencesAndOutputs(t *testing.T) { + ctx := t.Context() t.Parallel() pubmed := NewPubMedTool() @@ -268,7 +272,7 @@ func TestPubMed_ComponentReferencesAndOutputs(t *testing.T) { envelope := map[string]any{"results": []any{map[string]any{ "title": "Paper", "url": "https://pubmed.ncbi.nlm.nih.gov/1", "content": "Title: Paper\nAbstract: Evidence.", }}} - chunks, docAggs := pubmed.BuildReferences(context.Background(), envelope) + chunks, docAggs := pubmed.BuildReferences(ctx, envelope) if len(chunks) != 1 || len(docAggs) != 1 || chunks[0]["document_name"] != "Paper" { t.Fatalf("references = %#v / %#v", chunks, docAggs) } diff --git a/internal/agent/tool/qweather_test.go b/internal/agent/tool/qweather_test.go index c3e2bc824c..84e7990ce3 100644 --- a/internal/agent/tool/qweather_test.go +++ b/internal/agent/tool/qweather_test.go @@ -17,7 +17,6 @@ package tool import ( - "context" "encoding/json" "net/http" "net/http/httptest" @@ -88,6 +87,7 @@ func TestQWeather_BuildURL(t *testing.T) { } func TestQWeather_ParseResponse(t *testing.T) { + ctx := t.Context() t.Parallel() var gotQuery url.Values @@ -117,7 +117,7 @@ func TestQWeather_ParseResponse(t *testing.T) { }) tool := NewQWeatherToolWith(helper) - out, err := tool.InvokableRun(context.Background(), + out, err := tool.InvokableRun(ctx, `{"api_key":"K-abc","location":"101010100"}`) if err != nil { t.Fatalf("InvokableRun: %v", err) @@ -157,6 +157,7 @@ func TestQWeather_ParseResponse(t *testing.T) { } func TestQWeather_UpstreamBusinessError(t *testing.T) { + ctx := t.Context() t.Parallel() srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { @@ -174,7 +175,7 @@ func TestQWeather_UpstreamBusinessError(t *testing.T) { }) tool := NewQWeatherToolWith(helper) - _, err := tool.InvokableRun(context.Background(), + _, err := tool.InvokableRun(ctx, `{"api_key":"K-abc","location":"000000000"}`) if err == nil { t.Fatal("expected error for non-200 upstream code, got nil") @@ -185,10 +186,11 @@ func TestQWeather_UpstreamBusinessError(t *testing.T) { } func TestQWeather_RejectsMissingAPIKey(t *testing.T) { + ctx := t.Context() t.Parallel() tool := NewQWeatherTool() - _, err := tool.InvokableRun(context.Background(), + _, err := tool.InvokableRun(ctx, `{"location":"101010100"}`) if err == nil { t.Fatal("expected error for missing api_key") @@ -199,10 +201,11 @@ func TestQWeather_RejectsMissingAPIKey(t *testing.T) { } func TestQWeather_RejectsMissingLocation(t *testing.T) { + ctx := t.Context() t.Parallel() tool := NewQWeatherTool() - _, err := tool.InvokableRun(context.Background(), + _, err := tool.InvokableRun(ctx, `{"api_key":"K-abc"}`) if err == nil { t.Fatal("expected error for missing location") @@ -213,10 +216,11 @@ func TestQWeather_RejectsMissingLocation(t *testing.T) { } func TestQWeather_Info(t *testing.T) { + ctx := t.Context() t.Parallel() tool := NewQWeatherTool() - 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/retrieval_test.go b/internal/agent/tool/retrieval_test.go index 924ab1a258..3384d14194 100644 --- a/internal/agent/tool/retrieval_test.go +++ b/internal/agent/tool/retrieval_test.go @@ -29,10 +29,11 @@ import ( ) func TestRetrieval_StubsErrorWhenServiceMissing(t *testing.T) { + ctx := t.Context() t.Parallel() rt := NewRetrievalTool() - out, err := rt.InvokableRun(context.Background(), `{"query":"hello","dataset_ids":["kb-1"]}`) + out, err := rt.InvokableRun(ctx, `{"query":"hello","dataset_ids":["kb-1"]}`) if err == nil { t.Fatal("expected stub error, got nil") } @@ -54,9 +55,10 @@ func TestRetrieval_StubsErrorWhenServiceMissing(t *testing.T) { } func TestRetrieval_RejectsUseKG(t *testing.T) { + ctx := t.Context() t.Parallel() rt := NewRetrievalTool() - out, err := rt.InvokableRun(context.Background(), `{"query":"x","use_kg":true}`) + out, err := rt.InvokableRun(ctx, `{"query":"x","use_kg":true}`) if !errors.Is(err, ErrGraphRAGNotSupported) { t.Fatalf("err = %v, want ErrGraphRAGNotSupported", err) } @@ -66,10 +68,11 @@ func TestRetrieval_RejectsUseKG(t *testing.T) { } func TestRetrieval_InfoMatchesPythonMeta(t *testing.T) { + ctx := t.Context() t.Parallel() rt := NewRetrievalTool() - info, err := rt.Info(context.Background()) + info, err := rt.Info(ctx) if err != nil { t.Fatalf("Info: %v", err) } @@ -102,10 +105,11 @@ func TestRetrieval_InfoMatchesPythonMeta(t *testing.T) { } func TestRetrieval_EmptyArgsIsHandled(t *testing.T) { + ctx := t.Context() t.Parallel() rt := NewRetrievalTool() - out, err := rt.InvokableRun(context.Background(), "") + out, err := rt.InvokableRun(ctx, "") if err != nil { t.Fatalf("InvokableRun: %v", err) } @@ -126,7 +130,7 @@ func TestRetrieval_PassesTenantIDFromCanvasState(t *testing.T) { state := runtime.NewCanvasState("run-1", "task-1") state.Sys["tenant_id"] = "tenant-1" - ctx := runtime.WithState(context.Background(), state) + ctx := runtime.WithState(t.Context(), state) rt := NewRetrievalTool() _, err := rt.InvokableRun(ctx, `{"query":"hello","dataset_ids":["kb-1"]}`) @@ -146,7 +150,7 @@ func TestRetrieval_PassesUserIDWhenTenantIDMissing(t *testing.T) { state := runtime.NewCanvasState("run-1", "task-1") state.Sys["user_id"] = "user-1" - ctx := runtime.WithState(context.Background(), state) + ctx := runtime.WithState(t.Context(), state) rt := NewRetrievalTool() _, err := rt.InvokableRun(ctx, `{"query":"hello","dataset_ids":["kb-1"]}`) @@ -183,7 +187,9 @@ func TestRetrieval_UsesNodeParamsAsDefaults(t *testing.T) { t.Fatalf("BuildByName(retrieval) returned %T, want *RetrievalTool", built) } - _, err = rt.InvokableRun(context.Background(), `{"query":"hello"}`) + ctx := t.Context() + + _, err = rt.InvokableRun(ctx, `{"query":"hello"}`) if err != nil { t.Fatalf("InvokableRun: %v", err) } @@ -226,7 +232,8 @@ func TestRetrieval_ExplicitZeroSimilarityArgsOverrideDefaults(t *testing.T) { SimilarityThreshold: &similarityThreshold, KeywordsSimilarityWeight: &keywordsSimilarityWeight, }) - _, err := retrievalTool.InvokableRun(context.Background(), `{"query":"hello","similarity_threshold":0,"keywords_similarity_weight":0}`) + ctx := t.Context() + _, err := retrievalTool.InvokableRun(ctx, `{"query":"hello","similarity_threshold":0,"keywords_similarity_weight":0}`) if err != nil { t.Fatalf("InvokableRun: %v", err) } @@ -274,9 +281,10 @@ func TestRetrieval_ParsesEnhancementNodeParams(t *testing.T) { func TestRetrieval_UsesEmptyResponseForEmptyQuery(t *testing.T) { t.Parallel() + ctx := t.Context() rt := NewRetrievalToolWithDefaults(retrievalArgs{EmptyResponse: "No query or result."}) - out, err := rt.InvokableRun(context.Background(), `{"query":""}`) + out, err := rt.InvokableRun(ctx, `{"query":""}`) if err != nil { t.Fatalf("InvokableRun: %v", err) } @@ -314,7 +322,7 @@ func TestRetrieval_ResolvesCanvasVariables(t *testing.T) { state.SetVar("source", "ids", []any{"kb-1", "kb-2"}) state.SetVar("source", "query", "semantic question") state.SetVar("source", "value", "2026") - ctx := runtime.WithState(context.Background(), state) + ctx := runtime.WithState(t.Context(), state) ids, err := resolveRetrievalDatasetIDs(ctx, []string{"source@ids"}) if err != nil { @@ -347,17 +355,18 @@ func TestRetrieval_OmitsUnsetEmptyResponseFromArguments(t *testing.T) { } func TestRetrieval_UsesEmptyResponseWhenSearchHasNoChunks(t *testing.T) { + ctx := t.Context() prev := GetRetrievalService() SetRetrievalService(staticRetrievalService{}) t.Cleanup(func() { SetRetrievalService(prev) }) rt := NewRetrievalToolWithDefaults(retrievalArgs{DatasetIDs: []string{"kb-1"}, EmptyResponse: "No matching chunk."}) - out, err := rt.InvokableRun(context.Background(), `{"query":"love"}`) + out, err := rt.InvokableRun(ctx, `{"query":"love"}`) if err != nil { t.Fatalf("InvokableRun: %v", err) } var result retrievalResult - if err := json.Unmarshal([]byte(out), &result); err != nil { + if err = json.Unmarshal([]byte(out), &result); err != nil { t.Fatalf("unmarshal result: %v", err) } if result.FormalizedContent != "No matching chunk." { @@ -384,13 +393,14 @@ func TestRetrieval_IgnoresCanvasMetadataNodeParams(t *testing.T) { } func TestRetrieval_ModelArgsOverrideNodeDatasetIDs(t *testing.T) { + ctx := t.Context() prev := GetRetrievalService() svc := &capturingRetrievalService{} SetRetrievalService(svc) t.Cleanup(func() { SetRetrievalService(prev) }) rt := NewRetrievalToolWithDefaults(retrievalArgs{DatasetIDs: []string{"kb-default"}, TopN: 3}) - _, err := rt.InvokableRun(context.Background(), `{"query":"hello","dataset_ids":["kb-call"],"top_n":5}`) + _, err := rt.InvokableRun(ctx, `{"query":"hello","dataset_ids":["kb-call"],"top_n":5}`) if err != nil { t.Fatalf("InvokableRun: %v", err) } @@ -421,7 +431,7 @@ func TestRetrieval_RecordsFrontendReferencePayload(t *testing.T) { t.Cleanup(func() { SetRetrievalService(prev) }) state := runtime.NewCanvasState("run-1", "task-1") - ctx := runtime.WithState(context.Background(), state) + ctx := runtime.WithState(t.Context(), state) rt := NewRetrievalTool() _, err := rt.InvokableRun(ctx, `{"query":"hello","dataset_ids":["kb-1"]}`) diff --git a/internal/agent/tool/searxng_test.go b/internal/agent/tool/searxng_test.go index 22af5f3f0a..43b0755572 100644 --- a/internal/agent/tool/searxng_test.go +++ b/internal/agent/tool/searxng_test.go @@ -17,7 +17,6 @@ package tool import ( - "context" "encoding/json" "errors" "net" @@ -59,6 +58,7 @@ func TestSearXNGBuildURLMatchesPythonQuery(t *testing.T) { } func TestSearXNGInvokableRunPreservesRawResultsAndTopN(t *testing.T) { + ctx := t.Context() t.Parallel() server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, request *http.Request) { @@ -83,7 +83,7 @@ func TestSearXNGInvokableRunPreservesRawResultsAndTopN(t *testing.T) { defaults.SearXNGURL = server.URL defaults.TopN = 2 searchTool := newLocalSearXNGTool(t, defaults) - out, err := searchTool.InvokableRun(context.Background(), `{"query":" ragflow search ","searxng_url":"http://127.0.0.1:1"}`) + out, err := searchTool.InvokableRun(ctx, `{"query":" ragflow search ","searxng_url":"http://127.0.0.1:1"}`) if err != nil { t.Fatalf("InvokableRun: %v", err) } @@ -106,9 +106,10 @@ func TestSearXNGInvokableRunPreservesRawResultsAndTopN(t *testing.T) { } func TestSearXNGInfoMatchesPythonModelSchema(t *testing.T) { + ctx := t.Context() t.Parallel() - info, err := NewSearXNGTool().Info(context.Background()) + info, err := NewSearXNGTool().Info(ctx) if err != nil { t.Fatalf("Info: %v", err) } @@ -138,6 +139,7 @@ func TestSearXNGInfoMatchesPythonModelSchema(t *testing.T) { } func TestSearXNGEmptyTryRunInputsSkipRequest(t *testing.T) { + ctx := t.Context() t.Parallel() searchTool := NewSearXNGTool() @@ -151,12 +153,12 @@ func TestSearXNGEmptyTryRunInputsSkipRequest(t *testing.T) { `{"query":" ","searxng_url":"https://example.com"}`, `{"query":"ragflow"}`, } { - out, err := searchTool.InvokableRun(context.Background(), args) + out, err := searchTool.InvokableRun(ctx, args) if err != nil { t.Fatalf("InvokableRun(%s): %v", args, err) } var envelope searxngEnvelope - if err := json.Unmarshal([]byte(out), &envelope); err != nil || len(envelope.Results) != 0 || envelope.Error != "" { + if err = json.Unmarshal([]byte(out), &envelope); err != nil || len(envelope.Results) != 0 || envelope.Error != "" { t.Fatalf("InvokableRun(%s) = %s, want empty results", args, out) } } @@ -202,6 +204,7 @@ func TestSearXNGBuildByNameRejectsInvalidNodeParams(t *testing.T) { func TestSearXNGComponentContractReferencesAndOutputs(t *testing.T) { t.Parallel() + ctx := t.Context() tool := NewSearXNGTool() spec := tool.ComponentSpec() @@ -225,7 +228,7 @@ func TestSearXNGComponentContractReferencesAndOutputs(t *testing.T) { map[string]any{"title": "RAGFlow\nDocs", "url": "https://ragflow.io", "content": content, "engine": "bing", "score": 0.9}, map[string]any{"title": "Empty", "url": "https://example.com", "content": ""}, }} - chunks, docAggs := tool.BuildReferences(context.Background(), envelope) + chunks, docAggs := tool.BuildReferences(ctx, envelope) if len(chunks) != 1 || len(docAggs) != 1 { t.Fatalf("references = %#v / %#v", chunks, docAggs) } @@ -283,6 +286,7 @@ func TestRenderSearXNGReferencesStopsBeforeOverBudgetBlock(t *testing.T) { func TestSearXNGDoesNotRetryFailedRequest(t *testing.T) { t.Parallel() + ctx := t.Context() var calls atomic.Int32 server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { @@ -294,7 +298,7 @@ func TestSearXNGDoesNotRetryFailedRequest(t *testing.T) { defaults := defaultSearXNGParams() defaults.SearXNGURL = server.URL searchTool := newLocalSearXNGTool(t, defaults) - if _, err := searchTool.InvokableRun(context.Background(), `{"query":"single attempt"}`); err == nil { + if _, err := searchTool.InvokableRun(ctx, `{"query":"single attempt"}`); err == nil { t.Fatal("InvokableRun succeeded, want upstream error") } if calls.Load() != 1 { @@ -304,11 +308,12 @@ func TestSearXNGDoesNotRetryFailedRequest(t *testing.T) { func TestSearXNGSSRFGuardRejectsLoopback(t *testing.T) { t.Parallel() + ctx := t.Context() defaults := defaultSearXNGParams() defaults.SearXNGURL = "http://127.0.0.1:4000" searchTool := newSearXNGToolWithDefaults(nil, defaults) - out, err := searchTool.InvokableRun(context.Background(), `{"query":"metadata"}`) + out, err := searchTool.InvokableRun(ctx, `{"query":"metadata"}`) if err == nil || !errors.Is(err, ErrSSRFBlocked) { t.Fatalf("err = %v, want ErrSSRFBlocked", err) } diff --git a/internal/agent/tool/tavily_test.go b/internal/agent/tool/tavily_test.go index 72b082a339..987a36b805 100644 --- a/internal/agent/tool/tavily_test.go +++ b/internal/agent/tool/tavily_test.go @@ -17,7 +17,6 @@ package tool import ( - "context" "encoding/json" "net/http" "net/http/httptest" @@ -29,6 +28,7 @@ import ( func TestTavily_BuildRequest(t *testing.T) { t.Parallel() + ctx := t.Context() var gotPath, gotAuth, gotCT, gotMethod string var gotBody map[string]any @@ -44,14 +44,14 @@ func TestTavily_BuildRequest(t *testing.T) { })) defer srv.Close() - // Point the hard-coded tavily endpoint at the test server via a + // Point the hard-coded tavily endpoint at the test server via // transport that rewrites the host. Avoids the global-package-var // race that breaks parallel tests. helper := NewHTTPHelper().WithClient(&http.Client{ Transport: rewriteHostTransport(srv.URL), }) tool := NewTavilyToolWith(helper) - out, err := tool.InvokableRun(context.Background(), + out, err := tool.InvokableRun(ctx, `{"query":"ragflow","api_key":"key-xyz","max_results":3,"search_depth":"advanced","include_raw_content":true,"include_images":true}`) if err != nil { t.Fatalf("InvokableRun: %v", err) @@ -91,6 +91,7 @@ func TestTavily_BuildRequest(t *testing.T) { func TestTavily_ParseResponse(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", "application/json") @@ -108,7 +109,7 @@ func TestTavily_ParseResponse(t *testing.T) { Transport: rewriteHostTransport(srv.URL), }) tool := NewTavilyToolWith(helper) - out, err := tool.InvokableRun(context.Background(), + out, err := tool.InvokableRun(ctx, `{"query":"x","api_key":"k"}`) if err != nil { t.Fatalf("InvokableRun: %v", err) @@ -134,11 +135,12 @@ func TestTavily_ParseResponse(t *testing.T) { func TestTavily_RequiresAPIKey(t *testing.T) { t.Parallel() + ctx := t.Context() // envKey always returns "" so we know the failure is from the // missing api_key, not from a stray process env var. tool := NewTavilyToolWithEnvKey(NewHTTPHelper(), func() string { return "" }) - out, err := tool.InvokableRun(context.Background(), `{"query":"x"}`) + out, err := tool.InvokableRun(ctx, `{"query":"x"}`) if err != nil { t.Fatalf("InvokableRun should not return a Go error for missing api_key: %v", err) } @@ -149,6 +151,7 @@ func TestTavily_RequiresAPIKey(t *testing.T) { func TestTavily_APIKeyFromEnv(t *testing.T) { t.Parallel() + ctx := t.Context() var gotAuth string srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { @@ -161,7 +164,7 @@ func TestTavily_APIKeyFromEnv(t *testing.T) { Transport: rewriteHostTransport(srv.URL), }) tool := NewTavilyToolWithEnvKey(helper, func() string { return "from-env" }) - if _, err := tool.InvokableRun(context.Background(), `{"query":"x"}`); err != nil { + if _, err := tool.InvokableRun(ctx, `{"query":"x"}`); err != nil { t.Fatalf("InvokableRun: %v", err) } if gotAuth != "Bearer from-env" { @@ -171,9 +174,10 @@ func TestTavily_APIKeyFromEnv(t *testing.T) { func TestTavily_Info(t *testing.T) { t.Parallel() + ctx := t.Context() tool := NewTavilyTool() - info, err := tool.Info(context.Background()) + info, err := tool.Info(ctx) if err != nil { t.Fatalf("Info: %v", err) } @@ -187,15 +191,16 @@ func TestTavily_Info(t *testing.T) { func TestTavily_InfoArraySchemasIncludeStringItems(t *testing.T) { t.Parallel() + ctx := t.Context() - info, err := NewTavilyTool().Info(context.Background()) + info, err := NewTavilyTool().Info(ctx) if err != nil { t.Fatalf("Tavily Info: %v", err) } assertToolArrayItemType(t, info, "include_domains") assertToolArrayItemType(t, info, "exclude_domains") - extractInfo, err := NewTavilyExtractTool().Info(context.Background()) + extractInfo, err := NewTavilyExtractTool().Info(ctx) if err != nil { t.Fatalf("Tavily Extract Info: %v", err) } @@ -233,14 +238,15 @@ func assertToolArrayItemType(t *testing.T, info *schema.ToolInfo, fieldName stri func TestTavily_EmptyQueryReturnsEmptyResults(t *testing.T) { t.Parallel() + ctx := t.Context() tavily := NewTavilyToolWithEnvKey(NewHTTPHelper(), func() string { return "" }) - out, err := tavily.InvokableRun(context.Background(), `{"query":""}`) + out, err := tavily.InvokableRun(ctx, `{"query":""}`) if err != nil { t.Fatalf("InvokableRun(empty query): %v", err) } var envelope tavilyEnvelope - 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 != "" { @@ -250,6 +256,7 @@ func TestTavily_EmptyQueryReturnsEmptyResults(t *testing.T) { func TestTavily_PreservesRawResults(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", "application/json") @@ -257,12 +264,12 @@ func TestTavily_PreservesRawResults(t *testing.T) { })) defer srv.Close() helper := NewHTTPHelper().WithClient(&http.Client{Transport: rewriteHostTransport(srv.URL)}) - out, err := NewTavilyToolWith(helper).InvokableRun(context.Background(), `{"query":"x","api_key":"k"}`) + out, err := NewTavilyToolWith(helper).InvokableRun(ctx, `{"query":"x","api_key":"k"}`) if err != nil { t.Fatalf("InvokableRun: %v", err) } var envelope tavilyEnvelope - if err := json.Unmarshal([]byte(out), &envelope); err != nil { + if err = json.Unmarshal([]byte(out), &envelope); err != nil { t.Fatalf("decode response: %v", err) } custom, ok := envelope.Results[0]["custom"].(map[string]any) @@ -303,6 +310,7 @@ func TestTavily_BuildByNameUsesNodeDefaults(t *testing.T) { func TestTavily_ExplicitFlagsOverrideNodeDefaults(t *testing.T) { t.Parallel() + ctx := t.Context() var gotBody map[string]any srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { @@ -316,7 +324,7 @@ func TestTavily_ExplicitFlagsOverrideNodeDefaults(t *testing.T) { tavily := newTavilyTool(helper, func() string { return "" }, tavilyParams{ APIKey: "stored-key", IncludeAnswer: true, }) - if _, err := tavily.InvokableRun(context.Background(), `{"query":"ragflow"}`); err != nil { + if _, err := tavily.InvokableRun(ctx, `{"query":"ragflow"}`); err != nil { t.Fatalf("InvokableRun(node defaults): %v", err) } if gotBody["include_answer"] != true { @@ -325,7 +333,7 @@ func TestTavily_ExplicitFlagsOverrideNodeDefaults(t *testing.T) { if gotBody["include_raw_content"] != false || gotBody["include_images"] != false { t.Fatalf("include_raw_content/images should be forced false: %#v", gotBody) } - if _, err := tavily.InvokableRun(context.Background(), `{"query":"ragflow","include_answer":false}`); err != nil { + if _, err := tavily.InvokableRun(ctx, `{"query":"ragflow","include_answer":false}`); err != nil { t.Fatalf("InvokableRun: %v", err) } if gotBody["include_answer"] != false { @@ -354,6 +362,7 @@ func TestTavily_BuildByNameRejectsInvalidNodeDefaults(t *testing.T) { func TestTavily_ComponentReferencesAndOutputs(t *testing.T) { t.Parallel() + ctx := t.Context() tavily := NewTavilyTool() spec := tavily.ComponentSpec() @@ -373,7 +382,7 @@ func TestTavily_ComponentReferencesAndOutputs(t *testing.T) { "score": float64(0.75), "custom": "preserved", }}} - chunks, docAggs := tavily.BuildReferences(context.Background(), envelope) + chunks, docAggs := tavily.BuildReferences(ctx, envelope) if len(chunks) != 1 || len(docAggs) != 1 { t.Fatalf("references = %#v / %#v", chunks, docAggs) } @@ -398,6 +407,7 @@ func TestTavily_ComponentReferencesAndOutputs(t *testing.T) { func TestTavilyExtract_BuildRequest(t *testing.T) { t.Parallel() + ctx := t.Context() var gotPath, gotAuth, gotMethod string var gotBody map[string]any @@ -413,7 +423,7 @@ func TestTavilyExtract_BuildRequest(t *testing.T) { helper := NewHTTPHelper().WithClient(&http.Client{Transport: rewriteHostTransport(srv.URL)}) tool := NewTavilyExtractToolWith(helper) - out, err := tool.InvokableRun(context.Background(), + out, err := tool.InvokableRun(ctx, `{"urls":"https://a.example, https://b.example","api_key":"key-xyz","extract_depth":"advanced","format":"text"}`) if err != nil { t.Fatalf("InvokableRun: %v", err) @@ -444,6 +454,7 @@ func TestTavilyExtract_BuildRequest(t *testing.T) { func TestTavilyExtract_ParseResponse(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", "application/json") @@ -453,7 +464,7 @@ func TestTavilyExtract_ParseResponse(t *testing.T) { helper := NewHTTPHelper().WithClient(&http.Client{Transport: rewriteHostTransport(srv.URL)}) tool := NewTavilyExtractToolWith(helper) - out, err := tool.InvokableRun(context.Background(), `{"urls":["https://a.example/"],"api_key":"k"}`) + out, err := tool.InvokableRun(ctx, `{"urls":["https://a.example/"],"api_key":"k"}`) if err != nil { t.Fatalf("InvokableRun: %v", err) } @@ -478,9 +489,10 @@ func TestTavilyExtract_ParseResponse(t *testing.T) { func TestTavilyExtract_RequiresAPIKey(t *testing.T) { t.Parallel() + ctx := t.Context() tool := NewTavilyExtractToolWithEnvKey(NewHTTPHelper(), func() string { return "" }) - out, err := tool.InvokableRun(context.Background(), `{"urls":["https://a.example/"]}`) + out, err := tool.InvokableRun(ctx, `{"urls":["https://a.example/"]}`) if err != nil { t.Fatalf("InvokableRun should not return a Go error for missing api_key: %v", err) } @@ -491,9 +503,10 @@ func TestTavilyExtract_RequiresAPIKey(t *testing.T) { func TestTavilyExtract_Info(t *testing.T) { t.Parallel() + ctx := t.Context() tool := NewTavilyExtractTool() - info, err := tool.Info(context.Background()) + info, err := tool.Info(ctx) if err != nil { t.Fatalf("Info: %v", err) } @@ -571,6 +584,7 @@ func TestTavilyExtract_BuildByNameAcceptsNodeDefaults(t *testing.T) { func TestTavilyExtract_UsesNodeDefaults(t *testing.T) { t.Parallel() + ctx := t.Context() var gotAuth string var gotBody map[string]any @@ -589,7 +603,7 @@ func TestTavilyExtract_UsesNodeDefaults(t *testing.T) { ExtractDepth: "advanced", Format: "text", }) - if _, err := tavily.InvokableRun(context.Background(), `{"unrelated":"ignored"}`); err != nil { + if _, err := tavily.InvokableRun(ctx, `{"unrelated":"ignored"}`); err != nil { t.Fatalf("InvokableRun: %v", err) } if gotAuth != "Bearer stored-key" { diff --git a/internal/agent/tool/tushare_test.go b/internal/agent/tool/tushare_test.go index d7452b0780..6e2214b291 100644 --- a/internal/agent/tool/tushare_test.go +++ b/internal/agent/tool/tushare_test.go @@ -17,7 +17,6 @@ package tool import ( - "context" "encoding/json" "io" "net/http" @@ -88,6 +87,7 @@ func TestTushare_BuildRequest(t *testing.T) { func TestTushare_ParseResponse(t *testing.T) { t.Parallel() + ctx := t.Context() var ( gotMethod string @@ -122,7 +122,7 @@ func TestTushare_ParseResponse(t *testing.T) { }) tool := NewTushareToolWith(helper) - out, err := tool.InvokableRun(context.Background(), + out, err := tool.InvokableRun(ctx, `{"token":"T-test","api_name":"stock_basic","params":{"list_status":"L"}}`) if err != nil { t.Fatalf("InvokableRun: %v", err) @@ -169,9 +169,10 @@ func TestTushare_ParseResponse(t *testing.T) { func TestTushare_RejectsMissingToken(t *testing.T) { t.Parallel() + ctx := t.Context() tool := NewTushareTool() - _, err := tool.InvokableRun(context.Background(), + _, err := tool.InvokableRun(ctx, `{"api_name":"stock_basic"}`) if err == nil { t.Fatal("expected error for missing token") @@ -183,9 +184,10 @@ func TestTushare_RejectsMissingToken(t *testing.T) { func TestTushare_RejectsMissingAPIName(t *testing.T) { t.Parallel() + ctx := t.Context() tool := NewTushareTool() - _, err := tool.InvokableRun(context.Background(), + _, err := tool.InvokableRun(ctx, `{"token":"T-abc"}`) if err == nil { t.Fatal("expected error for missing api_name") @@ -197,6 +199,7 @@ func TestTushare_RejectsMissingAPIName(t *testing.T) { func TestTushare_UpstreamErrorCode(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", "application/json") @@ -213,7 +216,7 @@ func TestTushare_UpstreamErrorCode(t *testing.T) { }) tool := NewTushareToolWith(helper) - _, err := tool.InvokableRun(context.Background(), + _, err := tool.InvokableRun(ctx, `{"token":"T-abc","api_name":"premium_only"}`) if err == nil { t.Fatal("expected error for non-zero code, got nil") @@ -225,9 +228,10 @@ func TestTushare_UpstreamErrorCode(t *testing.T) { func TestTushare_Info(t *testing.T) { t.Parallel() + ctx := t.Context() tool := NewTushareTool() - 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/wencai_test.go b/internal/agent/tool/wencai_test.go index bc326baa80..7dcb07ea50 100644 --- a/internal/agent/tool/wencai_test.go +++ b/internal/agent/tool/wencai_test.go @@ -26,6 +26,7 @@ import ( func TestWencai_InvokeMatchesCurrentPythonResult(t *testing.T) { t.Parallel() + ctx := t.Context() cases := []struct { name string @@ -39,7 +40,7 @@ func TestWencai_InvokeMatchesCurrentPythonResult(t *testing.T) { for _, tc := range cases { t.Run(tc.name, func(t *testing.T) { t.Parallel() - out, err := NewWencaiTool().InvokableRun(context.Background(), tc.args) + out, err := NewWencaiTool().InvokableRun(ctx, tc.args) if err != nil { t.Fatalf("InvokableRun errored: %v (out=%s)", err, out) } @@ -59,8 +60,9 @@ func TestWencai_InvokeMatchesCurrentPythonResult(t *testing.T) { func TestWencai_RejectsMalformedJSON(t *testing.T) { t.Parallel() + ctx := t.Context() - out, err := NewWencaiTool().InvokableRun(context.Background(), `{not json`) + out, err := NewWencaiTool().InvokableRun(ctx, `{not json`) if err == nil { t.Fatal("expected malformed JSON error") } @@ -78,8 +80,9 @@ func TestWencai_RejectsMalformedJSON(t *testing.T) { func TestWencai_RespectsCanceledContext(t *testing.T) { t.Parallel() + ctx := t.Context() - ctx, cancel := context.WithCancel(context.Background()) + ctx, cancel := context.WithCancel(ctx) cancel() out, err := NewWencaiTool().InvokableRun(ctx, `{"query":"商业航天"}`) if !errors.Is(err, context.Canceled) { @@ -96,8 +99,9 @@ func TestWencai_RespectsCanceledContext(t *testing.T) { func TestWencai_InfoMatchesPythonMeta(t *testing.T) { t.Parallel() + ctx := t.Context() - info, err := NewWencaiTool().Info(context.Background()) + info, err := NewWencaiTool().Info(ctx) if err != nil { t.Fatalf("Info: %v", err) } diff --git a/internal/agent/tool/wikipedia_test.go b/internal/agent/tool/wikipedia_test.go index e29cf03100..0a140e2c62 100644 --- a/internal/agent/tool/wikipedia_test.go +++ b/internal/agent/tool/wikipedia_test.go @@ -17,7 +17,6 @@ package tool import ( - "context" "encoding/json" "net/http" "net/http/httptest" @@ -97,6 +96,7 @@ func TestWikipedia_BuildURL(t *testing.T) { func TestWikipedia_ParseResults(t *testing.T) { t.Parallel() + ctx := t.Context() var gotUA string srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { @@ -114,12 +114,12 @@ func TestWikipedia_ParseResults(t *testing.T) { defer srv.Close() // Point the hard-coded en.wikipedia.org endpoint at the test server - // by injecting a transport that rewrites the request host. + // by injecting transport that rewrites the request host. helper := NewHTTPHelper().WithClient(&http.Client{ Transport: rewriteHostTransport(srv.URL), }) tool := NewWikipediaToolWith(helper) - out, err := tool.InvokableRun(context.Background(), `{"query":"RAG","lang":"en","max_results":5}`) + out, err := tool.InvokableRun(ctx, `{"query":"RAG","lang":"en","max_results":5}`) if err != nil { t.Fatalf("InvokableRun: %v", err) } @@ -179,9 +179,10 @@ func (t *hostSwapRT) RoundTrip(req *http.Request) (*http.Response, error) { func TestWikipedia_Info(t *testing.T) { t.Parallel() + ctx := t.Context() tool := NewWikipediaTool() - info, err := tool.Info(context.Background()) + info, err := tool.Info(ctx) if err != nil { t.Fatalf("Info: %v", err) } @@ -195,9 +196,10 @@ func TestWikipedia_Info(t *testing.T) { func TestWikipedia_EmptyQuery(t *testing.T) { t.Parallel() + ctx := t.Context() tool := NewWikipediaTool() - out, err := tool.InvokableRun(context.Background(), `{"query":""}`) + out, err := tool.InvokableRun(ctx, `{"query":""}`) if err != nil { t.Fatalf("InvokableRun(empty): %v", err) } @@ -212,6 +214,7 @@ func TestWikipedia_EmptyQuery(t *testing.T) { func TestWikipedia_ComponentReferencesAndOutputs(t *testing.T) { t.Parallel() + ctx := t.Context() wikipedia := NewWikipediaTool() spec := wikipedia.ComponentSpec() @@ -223,7 +226,7 @@ func TestWikipedia_ComponentReferencesAndOutputs(t *testing.T) { "url": "https://en.wikipedia.org/wiki/RAG", "content": "RAG is an acronym.", }}} - chunks, docAggs := wikipedia.BuildReferences(context.Background(), envelope) + chunks, docAggs := wikipedia.BuildReferences(ctx, envelope) if len(chunks) != 1 || len(docAggs) != 1 || chunks[0]["document_name"] != "RAG" || chunks[0]["similarity"] != 1 { t.Fatalf("references = %#v / %#v", chunks, docAggs) } diff --git a/internal/agent/tool/yahoo_finance_test.go b/internal/agent/tool/yahoo_finance_test.go index 6196be870d..b7c9b84fca 100644 --- a/internal/agent/tool/yahoo_finance_test.go +++ b/internal/agent/tool/yahoo_finance_test.go @@ -17,7 +17,6 @@ package tool import ( - "context" "encoding/json" "net/http" "net/http/httptest" @@ -50,6 +49,7 @@ func TestYahooFinanceBuildURL(t *testing.T) { func TestYahooFinanceInvokableRunBuildsMarkdownReport(t *testing.T) { t.Parallel() + ctx := t.Context() var gotQuery, gotUserAgent string var gotPaths []string @@ -92,7 +92,7 @@ func TestYahooFinanceInvokableRunBuildsMarkdownReport(t *testing.T) { defer server.Close() helper := NewHTTPHelper().WithClient(&http.Client{Transport: rewriteHostTransport(server.URL)}) - raw, err := NewYahooFinanceToolWith(helper).InvokableRun(context.Background(), `{"stock_code":" AAPL "}`) + raw, err := NewYahooFinanceToolWith(helper).InvokableRun(ctx, `{"stock_code":" AAPL "}`) if err != nil { t.Fatalf("InvokableRun: %v", err) } @@ -132,6 +132,7 @@ func TestYahooFinanceInvokableRunBuildsMarkdownReport(t *testing.T) { func TestYahooFinanceInvokableRunAcceptsQueryAlias(t *testing.T) { t.Parallel() + ctx := t.Context() var gotQuery string server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { @@ -158,7 +159,7 @@ func TestYahooFinanceInvokableRunAcceptsQueryAlias(t *testing.T) { defer server.Close() helper := NewHTTPHelper().WithClient(&http.Client{Transport: rewriteHostTransport(server.URL)}) - raw, err := NewYahooFinanceToolWith(helper).InvokableRun(context.Background(), `{"query":"3800.HK"}`) + raw, err := NewYahooFinanceToolWith(helper).InvokableRun(ctx, `{"query":"3800.HK"}`) if err != nil { t.Fatalf("InvokableRun: %v", err) } @@ -166,7 +167,7 @@ func TestYahooFinanceInvokableRunAcceptsQueryAlias(t *testing.T) { t.Fatalf("q = %q", gotQuery) } var envelope yahooFinanceEnvelope - if err := json.Unmarshal([]byte(raw), &envelope); err != nil { + if err = json.Unmarshal([]byte(raw), &envelope); err != nil { t.Fatalf("unmarshal output %q: %v", raw, err) } if !strings.Contains(envelope.Report, "| symbol | 3800.HK |") || @@ -177,6 +178,7 @@ func TestYahooFinanceInvokableRunAcceptsQueryAlias(t *testing.T) { func TestYahooFinanceUsesSearchResolvedSymbolForDownstreamRequests(t *testing.T) { t.Parallel() + ctx := t.Context() var paths []string var gotSearchQuery string @@ -232,7 +234,7 @@ func TestYahooFinanceUsesSearchResolvedSymbolForDownstreamRequests(t *testing.T) helper := NewHTTPHelper().WithClient(&http.Client{Transport: rewriteHostTransport(server.URL)}) yahoo := NewYahooFinanceToolWithDefaults(helper, yahooFinanceParams{Info: true, Count: true}) - raw, err := yahoo.InvokableRun(context.Background(), `{"stock_code":"Apple"}`) + raw, err := yahoo.InvokableRun(ctx, `{"stock_code":"Apple"}`) if err != nil { t.Fatalf("InvokableRun: %v", err) } @@ -249,7 +251,7 @@ func TestYahooFinanceUsesSearchResolvedSymbolForDownstreamRequests(t *testing.T) } } var envelope yahooFinanceEnvelope - if err := json.Unmarshal([]byte(raw), &envelope); err != nil { + if err = json.Unmarshal([]byte(raw), &envelope); err != nil { t.Fatalf("unmarshal output %q: %v", raw, err) } if !strings.Contains(envelope.Report, "| longname | Apple Inc. |") || @@ -260,6 +262,7 @@ func TestYahooFinanceUsesSearchResolvedSymbolForDownstreamRequests(t *testing.T) func TestYahooFinanceEmptyStockCodeSkipsRequest(t *testing.T) { t.Parallel() + ctx := t.Context() calls := 0 server := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) { @@ -269,12 +272,12 @@ func TestYahooFinanceEmptyStockCodeSkipsRequest(t *testing.T) { helper := NewHTTPHelper().WithClient(&http.Client{Transport: rewriteHostTransport(server.URL)}) for _, args := range []string{`{"stock_code":""}`, `{"stock_code":" "}`} { - raw, err := NewYahooFinanceToolWith(helper).InvokableRun(context.Background(), args) + raw, err := NewYahooFinanceToolWith(helper).InvokableRun(ctx, args) if err != nil { t.Fatalf("InvokableRun(%s): %v", args, err) } var envelope yahooFinanceEnvelope - if err := json.Unmarshal([]byte(raw), &envelope); err != nil { + if err = json.Unmarshal([]byte(raw), &envelope); err != nil { t.Fatalf("unmarshal output: %v", err) } if envelope.Report != "" || envelope.Error != "" { @@ -288,6 +291,7 @@ func TestYahooFinanceEmptyStockCodeSkipsRequest(t *testing.T) { func TestYahooFinanceAllSectionsDisabledSkipsRequest(t *testing.T) { t.Parallel() + ctx := t.Context() calls := 0 server := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) { @@ -297,12 +301,12 @@ func TestYahooFinanceAllSectionsDisabledSkipsRequest(t *testing.T) { helper := NewHTTPHelper().WithClient(&http.Client{Transport: rewriteHostTransport(server.URL)}) yahoo := NewYahooFinanceToolWithDefaults(helper, yahooFinanceParams{}) - raw, err := yahoo.InvokableRun(context.Background(), `{"stock_code":"AAPL"}`) + raw, err := yahoo.InvokableRun(ctx, `{"stock_code":"AAPL"}`) if err != nil { t.Fatalf("InvokableRun: %v", err) } var envelope yahooFinanceEnvelope - if err := json.Unmarshal([]byte(raw), &envelope); err != nil { + if err = json.Unmarshal([]byte(raw), &envelope); err != nil { t.Fatalf("unmarshal output: %v", err) } if envelope.Report != "" || calls != 0 { @@ -312,6 +316,7 @@ func TestYahooFinanceAllSectionsDisabledSkipsRequest(t *testing.T) { func TestYahooFinanceSummaryUsesCrumbAndCookie(t *testing.T) { t.Parallel() + ctx := t.Context() var sawCookieOnCrumb, sawCookieOnSummary bool server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { @@ -352,7 +357,7 @@ func TestYahooFinanceSummaryUsesCrumbAndCookie(t *testing.T) { helper := NewHTTPHelper().WithClient(&http.Client{Transport: rewriteHostTransport(server.URL)}) yahoo := NewYahooFinanceToolWithDefaults(helper, yahooFinanceParams{Count: true}) - raw, err := yahoo.InvokableRun(context.Background(), `{"stock_code":"AAPL"}`) + raw, err := yahoo.InvokableRun(ctx, `{"stock_code":"AAPL"}`) if err != nil { t.Fatalf("InvokableRun: %v", err) } @@ -360,7 +365,7 @@ func TestYahooFinanceSummaryUsesCrumbAndCookie(t *testing.T) { t.Fatalf("cookie propagation crumb=%v summary=%v", sawCookieOnCrumb, sawCookieOnSummary) } var envelope yahooFinanceEnvelope - if err := json.Unmarshal([]byte(raw), &envelope); err != nil { + if err = json.Unmarshal([]byte(raw), &envelope); err != nil { t.Fatalf("unmarshal output: %v", err) } if !strings.Contains(envelope.Report, "# Count:") || @@ -459,6 +464,7 @@ func TestMergeYahooFinanceParamsKeepsStockCodeAndUsesNodeConfig(t *testing.T) { func TestYahooFinanceErrorsReturnEnvelope(t *testing.T) { t.Parallel() + ctx := t.Context() tests := []struct { name string @@ -484,7 +490,7 @@ func TestYahooFinanceErrorsReturnEnvelope(t *testing.T) { defer server.Close() helper := NewHTTPHelperWithRetry(RetryConfig{MaxAttempts: 1}).WithClient(&http.Client{Transport: rewriteHostTransport(server.URL)}) - raw, err := NewYahooFinanceToolWith(helper).InvokableRun(context.Background(), `{"stock_code":"AAPL"}`) + raw, err := NewYahooFinanceToolWith(helper).InvokableRun(ctx, `{"stock_code":"AAPL"}`) if err == nil || !strings.Contains(err.Error(), test.wantError) { t.Fatalf("err = %v, want %q", err, test.wantError) } @@ -501,8 +507,9 @@ func TestYahooFinanceErrorsReturnEnvelope(t *testing.T) { func TestYahooFinanceMalformedArguments(t *testing.T) { t.Parallel() + ctx := t.Context() - raw, err := NewYahooFinanceTool().InvokableRun(context.Background(), `{`) + raw, err := NewYahooFinanceTool().InvokableRun(ctx, `{`) if err == nil || !strings.Contains(err.Error(), "parse arguments") { t.Fatalf("err = %v", err) } @@ -535,8 +542,9 @@ func TestYahooFinanceComponentContract(t *testing.T) { func TestYahooFinanceInfoExposesPythonCompatibleParams(t *testing.T) { t.Parallel() + ctx := t.Context() - info, err := NewYahooFinanceTool().Info(context.Background()) + info, err := NewYahooFinanceTool().Info(ctx) if err != nil { t.Fatalf("Info: %v", err) }