133 lines
3.9 KiB
JavaScript
133 lines
3.9 KiB
JavaScript
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("==========================================");
|