Luigit
repositories / pi-ext

pi-ext

bugabingas pi extensions

owned by admin

extensions/web/__tests__/search.test.ts

Raw
import { afterEach, describe, expect, it, vi } from "vitest";
import * as codex from "../codex";
import * as ddg from "../ddg";
import * as fetchClient from "../fetch";
import * as google from "../google";
import * as minimax from "../minimax";
import * as parallel from "../parallel";
import { searchWeb } from "../search";
import * as zai from "../zai";

afterEach(() => vi.restoreAllMocks());

describe("search provider fallback order", () => {
	it.each(["google", "parallel", "codex", "minimax", "ddg", "zai", "none"])(
		"tries the configured order, with Z.AI last, until %s succeeds",
		async (winner) => {
			const calls: string[] = [];
			const attempt = (name: string) => {
				calls.push(name);
				if (name !== winner) throw new Error(`${name} unavailable`);
			};
			vi.spyOn(codex, "isCodexAvailable").mockResolvedValue(true);
			vi.spyOn(minimax, "isMinimaxAvailable").mockResolvedValue(true);
			vi.spyOn(zai, "isZaiAvailable").mockResolvedValue(true);
			vi.spyOn(ddg, "throttleWait").mockReturnValue(0);
			vi.spyOn(ddg, "markSearchTime").mockImplementation(() => {});
			vi.spyOn(codex, "codexSearch").mockImplementation(async () => {
				attempt("codex");
				return [];
			});
			vi.spyOn(minimax, "minimaxSearch").mockImplementation(async () => {
				attempt("minimax");
				return {
					organic: [],
					related_searches: [],
					base_resp: { status_code: 0, status_msg: "ok" },
				};
			});
			vi.spyOn(parallel, "parallelSearch").mockImplementation(async () => {
				attempt("parallel");
				return [];
			});
			vi.spyOn(google, "googleSearch").mockImplementation(async () => {
				attempt("google");
				return [];
			});
			vi.spyOn(fetchClient, "fetchSearchHtml").mockImplementation(async () => {
				attempt("ddg");
				return "";
			});
			vi.spyOn(ddg, "parseDdgResults").mockReturnValue({
				results: [],
				diagnostics: { resultsFound: 0, resultsSkipped: 0, rawBlockCount: 0 },
			});
			vi.spyOn(zai, "zaiSearch").mockImplementation(async () => {
				attempt("zai");
				return [];
			});

			const result = searchWeb({
				query: "docs",
				count: 5,
				timeout: 5000,
				ctx: {} as any,
			});
			const order = ["google", "parallel", "codex", "minimax", "ddg", "zai"];
			if (winner === "none") {
				await expect(result).rejects.toThrow(
					"Search unavailable: Google WML: google unavailable | Parallel: parallel unavailable | Codex: codex unavailable | MiniMax: minimax unavailable | DDG: ddg unavailable | Z.AI: zai unavailable",
				);
				expect(calls).toEqual(order);
			} else {
				await expect(result).resolves.toEqual({ results: [], source: winner });
				expect(calls).toEqual(order.slice(0, order.indexOf(winner) + 1));
			}
		},
	);

	it("does not continue to another provider after cancellation", async () => {
		vi.spyOn(codex, "isCodexAvailable").mockResolvedValue(true);
		vi.spyOn(minimax, "isMinimaxAvailable").mockResolvedValue(false);
		vi.spyOn(zai, "isZaiAvailable").mockResolvedValue(false);
		const abort = new AbortController();
		vi.spyOn(google, "googleSearch").mockImplementation(async () => {
			abort.abort(new Error("cancelled"));
			throw abort.signal.reason;
		});
		const next = vi.spyOn(parallel, "parallelSearch");
		await expect(
			searchWeb({
				query: "docs",
				count: 5,
				timeout: 5000,
				signal: abort.signal,
				ctx: {} as any,
			}),
		).rejects.toThrow("cancelled");
		expect(next).not.toHaveBeenCalled();
	});
});