Files
diegosouzapw--omniroute/tests/unit/embeddings-proxy-forwarding.test.ts
2026-07-13 13:39:12 +08:00

111 lines
4.2 KiB
TypeScript

import test from "node:test";
import assert from "node:assert/strict";
import fs from "node:fs";
import os from "node:os";
import path from "node:path";
import http from "node:http";
// Isolate the DB to a temp dir BEFORE importing any module that opens it.
const TEST_DATA_DIR = fs.mkdtempSync(path.join(os.tmpdir(), "omniroute-embed-proxy-"));
process.env.DATA_DIR = TEST_DATA_DIR;
const core = await import("../../src/lib/db/core.ts");
const providersDb = await import("../../src/lib/db/providers.ts");
const settingsDb = await import("../../src/lib/db/settings.ts");
const { createEmbeddingResponse } = await import("../../src/lib/embeddings/service.ts");
const { resolveProxyForRequest } = await import("../../open-sse/utils/proxyFetch.ts");
test.after(() => {
core.resetDbInstance();
fs.rmSync(TEST_DATA_DIR, { recursive: true, force: true });
});
async function withHttpServer(handler: http.RequestListener, fn: (baseUrl: string) => Promise<void>) {
const server = http.createServer(handler);
await new Promise<void>((resolve, reject) => {
server.once("error", reject);
server.listen(0, "127.0.0.1", () => resolve());
});
const address = server.address();
assert.ok(address && typeof address === "object");
try {
await fn(`http://127.0.0.1:${(address as { port: number }).port}`);
} finally {
await new Promise<void>((resolve, reject) => {
server.close((err) => (err ? reject(err) : resolve()));
});
}
}
test("embeddings forward the connection-level (key) pinned proxy to the upstream fetch", async () => {
// T14 Proxy Fast-Fail performs a real TCP reachability check before running
// the request inside the proxy context, so the "proxy" must be a real,
// reachable local listener (its actual response body is irrelevant here —
// the assertion is about which proxy context the upstream fetch observes).
await withHttpServer(
(_req, res) => {
res.writeHead(200);
res.end("ok");
},
async (proxyBaseUrl) => {
const proxyUrl = new URL(proxyBaseUrl);
const connection = await providersDb.createProviderConnection({
provider: "mistral",
authType: "apikey",
name: "Test Mistral Proxy",
apiKey: "mistral-test-key",
});
// Pin a proxy at the connection ("key") level — the most specific level,
// exactly like a user would configure per-connection in the dashboard.
await settingsDb.setProxyForLevel("key", (connection as any).id, {
type: "http",
host: proxyUrl.hostname,
port: Number(proxyUrl.port),
});
let capturedProxySource: string | null = null;
let capturedProxyUrl: string | null = null;
const originalFetch = globalThis.fetch;
globalThis.fetch = async (input: RequestInfo | URL) => {
const targetUrl = typeof input === "string" ? input : (input as URL).toString();
const resolved = resolveProxyForRequest(targetUrl);
capturedProxySource = resolved.source;
capturedProxyUrl = resolved.proxyUrl;
return new Response(
JSON.stringify({
data: [{ object: "embedding", embedding: [0.1, 0.2], index: 0 }],
usage: { prompt_tokens: 3, total_tokens: 3 },
}),
{ status: 200, headers: { "content-type": "application/json" } }
);
};
try {
const res = await createEmbeddingResponse({
model: "mistral/mistral-embed",
input: "hello world",
});
assert.equal(res.status, 200, "embedding request should succeed");
} finally {
globalThis.fetch = originalFetch;
}
// The upstream fetch must have run inside the AsyncLocalStorage proxy
// context carrying the connection's pinned proxy — not "direct" (no
// context) and not the env-var proxy fallback.
assert.equal(
capturedProxySource,
"context",
`expected the embeddings upstream fetch to run inside a proxy context, got source="${capturedProxySource}"`
);
assert.ok(
capturedProxyUrl && capturedProxyUrl.includes(`${proxyUrl.hostname}:${proxyUrl.port}`),
`expected the connection's pinned proxy to be forwarded, got "${capturedProxyUrl}"`
);
}
);
});