From aad4a0daee798be8f0c649d40ffbee27682e93a0 Mon Sep 17 00:00:00 2001 From: WeiCHEN Date: Mon, 7 Sep 2026 17:31:35 +0800 Subject: [PATCH] test(agent): add Tavily pool rotation and PiAgent remote runner tests Co-authored-by: Cursor --- deploy/pi-agent/tests/tavily-test.mjs | 132 ++++++++++++++++++++++++++ internal/agent/backend_test.go | 61 ++++++++++++ 2 files changed, 193 insertions(+) create mode 100644 deploy/pi-agent/tests/tavily-test.mjs diff --git a/deploy/pi-agent/tests/tavily-test.mjs b/deploy/pi-agent/tests/tavily-test.mjs new file mode 100644 index 00000000..61f77db4 --- /dev/null +++ b/deploy/pi-agent/tests/tavily-test.mjs @@ -0,0 +1,132 @@ +import assert from "node:assert/strict"; +import { + TavilyKeyPool, + formatTavilyResults, + default as registerTavilyExtension, +} from "../extensions/tavily-search.mjs"; + +console.log("▶ Testing TavilyKeyPool initialization and parsing..."); + +// 1. Test key parsing from multiple env formats +{ + const pool = TavilyKeyPool.fromEnv({ + TAVILY_API_KEYS: "tvly-key1, tvly-key2; tvly-key3\ntvly-key4", + }); + assert.equal(pool.size(), 4, "Should parse 4 keys across comma, semicolon, space, newline"); + assert.equal(pool.getCurrentKey(), "tvly-key1"); +} + +{ + const poolSingle = TavilyKeyPool.fromEnv({ + TAVILY_API_KEY: "tvly-single-key-12345", + }); + assert.equal(poolSingle.size(), 1); + assert.equal(poolSingle.getCurrentKey(), "tvly-single-key-12345"); +} + +{ + const poolEmpty = TavilyKeyPool.fromEnv({}); + assert.equal(poolEmpty.size(), 0); + assert.equal(poolEmpty.getCurrentKey(), null); +} + +console.log("✔ Key parsing tests passed."); + +// 2. Test automatic failover and key rotation +console.log("▶ Testing automatic failover & rotation in pool..."); +{ + const keys = ["tvly-bad-key-1", "tvly-rate-limited-2", "tvly-good-key-3"]; + const pool = new TavilyKeyPool(keys); + + const attemptedKeys = []; + const result = await pool.executeWithPool(async (key) => { + attemptedKeys.push(key); + if (key === "tvly-bad-key-1") { + const err = new Error("Invalid API key"); + err.status = 401; + throw err; + } + if (key === "tvly-rate-limited-2") { + const err = new Error("Usage limit reached"); + err.status = 429; + throw err; + } + return { status: "success", data: "live search results" }; + }); + + assert.deepEqual(attemptedKeys, [ + "tvly-bad-key-1", + "tvly-rate-limited-2", + "tvly-good-key-3", + ]); + assert.equal(result.status, "success"); + console.log("✔ Auto-rotation to working key passed."); +} + +// 3. Test behavior when all keys fail +console.log("▶ Testing exhaustion when all keys in pool fail..."); +{ + const keys = ["key-fail-1", "key-fail-2"]; + const pool = new TavilyKeyPool(keys); + + let caughtError = null; + try { + await pool.executeWithPool(async (key) => { + const err = new Error(`Quota exceeded for ${key}`); + err.status = 429; + throw err; + }); + } catch (err) { + caughtError = err; + } + + assert.ok(caughtError !== null, "Should throw when all keys fail"); + assert.match(caughtError.message, /All 2 Tavily API key\(s\) in pool failed/); + console.log("✔ Exhaustion error handling passed."); +} + +// 4. Test formatTavilyResults +console.log("▶ Testing formatTavilyResults markdown generation..."); +{ + const mockData = { + answer: "Medical robotics includes surgical robots like da Vinci and orthopedic arms.", + results: [ + { + title: "Medical Robotics 2026 Overview", + url: "https://example.com/robotics", + content: "Detailed overview of surgical robotics developments in 2026.", + score: 0.985, + }, + ], + }; + + const md = formatTavilyResults(mockData, "medical robotics"); + assert.match(md, /Tavily 实时网络搜索结果/); + assert.match(md, /Medical robotics includes surgical robots/); + assert.match(md, /https:\/\/example.com\/robotics/); + console.log("✔ Markdown formatting passed."); +} + +// 5. Test tool registration on Pi mock +console.log("▶ Testing Pi extension registration..."); +{ + const registeredTools = []; + const mockPi = { + registerTool(tool) { + registeredTools.push(tool); + }, + getAllTools() { + return registeredTools; + }, + }; + + registerTavilyExtension(mockPi); + assert.equal(registeredTools.length, 2, "Should register tavily_search and web_search alias"); + assert.equal(registeredTools[0].name, "tavily_search"); + assert.equal(registeredTools[1].name, "web_search"); + console.log("✔ Tool registration on Pi passed."); +} + +console.log("=========================================="); +console.log("🎉 ALL TAVILY TESTS PASSED SUCCESSFULLY!"); +console.log("=========================================="); diff --git a/internal/agent/backend_test.go b/internal/agent/backend_test.go index 3a4834b9..ffaced36 100644 --- a/internal/agent/backend_test.go +++ b/internal/agent/backend_test.go @@ -2,6 +2,9 @@ package agent_test import ( "context" + "encoding/json" + "net/http" + "net/http/httptest" "os" "path/filepath" "testing" @@ -40,3 +43,61 @@ func TestStubBackendExecution(t *testing.T) { t.Fatalf("expected output.md to exist: %v", err) } } + +func TestPiAgentRemoteRunner(t *testing.T) { + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/v1/run" { + http.NotFound(w, r) + return + } + var payload map[string]any + _ = json.NewDecoder(r.Body).Decode(&payload) + + resp := map[string]any{ + "ok": true, + "patches": []map[string]string{ + { + "rel_path": "output.md", + "content": "# Remote Generated Deliverable\nValidated from Container Runner", + }, + }, + } + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(resp) + })) + defer ts.Close() + + tmpDir := t.TempDir() + cfg := &config.Config{ + AgentBackend: "pi-agent", + PiAgentRunnerURL: ts.URL, + PiAgentModel: "MiniMax-M3", + PiAgentProvider: "minimax", + AgentTimeout: 5 * time.Second, + } + + backend := agent.NewBackend("pi-agent", cfg) + res, err := backend.Execute(context.Background(), agent.Request{ + WorkDir: tmpDir, + Prompt: "Generate remote test document", + Targets: []string{"output.md"}, + }) + if err != nil { + t.Fatalf("remote execute returned unexpected error: %v", err) + } + if !res.Success { + t.Fatalf("expected success, got error: %s", res.ErrorReason) + } + if string(res.Patches["output.md"]) != "# Remote Generated Deliverable\nValidated from Container Runner" { + t.Errorf("unexpected content: %s", string(res.Patches["output.md"])) + } + + // Verify local file written + localContent, err := os.ReadFile(filepath.Join(tmpDir, "output.md")) + if err != nil { + t.Fatalf("expected local file written: %v", err) + } + if string(localContent) != "# Remote Generated Deliverable\nValidated from Container Runner" { + t.Errorf("unexpected local file content: %s", string(localContent)) + } +}