feat: add extensive test suite
This commit is contained in:
@@ -0,0 +1,165 @@
|
||||
import { describe, expect, mock, test } from "bun:test";
|
||||
import { createTtsHandler, pcmToWav } from "./handler";
|
||||
|
||||
const request = (message = "hello", key = "widget") =>
|
||||
new Request(
|
||||
`https://example.com/api/tts?message=${encodeURIComponent(message)}&key=${encodeURIComponent(key)}`,
|
||||
);
|
||||
|
||||
function dependencies() {
|
||||
const redis = {
|
||||
get: mock((): Promise<string | null> => Promise.resolve(null)),
|
||||
setex: mock(() => Promise.resolve("OK")),
|
||||
};
|
||||
return {
|
||||
redis,
|
||||
widgetKey: "widget",
|
||||
apiKeys: "first, second",
|
||||
generateContent: mock(
|
||||
(_options: {
|
||||
apiKey: string;
|
||||
model: string;
|
||||
message: string;
|
||||
voice: string;
|
||||
}): Promise<string | undefined> => Promise.resolve("generated-pcm"),
|
||||
),
|
||||
hash: () => "hash",
|
||||
random: () => 0,
|
||||
queue: (callback: () => void) => callback(),
|
||||
toWav: (audio: string) => `wav:${audio}`,
|
||||
log: mock(() => {}),
|
||||
logError: mock(() => {}),
|
||||
};
|
||||
}
|
||||
|
||||
describe("TTS handler", () => {
|
||||
test.each([
|
||||
{ widgetKey: undefined, apiKeys: "key" },
|
||||
{ widgetKey: "widget", apiKeys: undefined },
|
||||
{ widgetKey: "widget", apiKeys: " , " },
|
||||
])("returns 503 when configuration is missing", async (config) => {
|
||||
const deps = { ...dependencies(), ...config };
|
||||
const response = await createTtsHandler(deps)(request());
|
||||
expect(response.status).toBe(503);
|
||||
expect(await response.json()).toEqual({ error: "TTS is not configured" });
|
||||
expect(deps.redis.get).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
test("rejects an invalid widget key before accessing Redis", async () => {
|
||||
const deps = dependencies();
|
||||
const response = await createTtsHandler(deps)(request("hello", "wrong"));
|
||||
expect(response.status).toBe(400);
|
||||
expect(await response.text()).toBe("Invalid key");
|
||||
expect(deps.redis.get).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
test("rejects messages over 1000 characters", async () => {
|
||||
const deps = dependencies();
|
||||
const response = await createTtsHandler(deps)(request("x".repeat(1001)));
|
||||
expect(response.status).toBe(400);
|
||||
expect(deps.generateContent).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
test("returns cached audio without generating content", async () => {
|
||||
const deps = dependencies();
|
||||
deps.redis.get.mockResolvedValueOnce("cached-pcm");
|
||||
const response = await createTtsHandler(deps)(request());
|
||||
expect(response.status).toBe(200);
|
||||
expect(response.headers.get("content-type")).toBe("audio/wav");
|
||||
expect(await response.text()).toBe("wav:cached-pcm");
|
||||
expect(deps.redis.get).toHaveBeenCalledWith("tts:hash");
|
||||
expect(deps.generateContent).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
test("uses the persisted rotation, selected voice, and caches generated audio", async () => {
|
||||
const deps = dependencies();
|
||||
deps.redis.get.mockResolvedValueOnce(null).mockResolvedValueOnce("1");
|
||||
const response = await createTtsHandler(deps)(request("speak"));
|
||||
|
||||
expect(await response.text()).toBe("wav:generated-pcm");
|
||||
expect(deps.generateContent).toHaveBeenCalledWith({
|
||||
apiKey: "second",
|
||||
model: "gemini-2.5-flash-preview-tts",
|
||||
message: "speak",
|
||||
voice: "Zephyr",
|
||||
});
|
||||
expect(deps.redis.setex).toHaveBeenCalledWith(
|
||||
"tts:hash",
|
||||
604800,
|
||||
"generated-pcm",
|
||||
);
|
||||
expect(deps.redis.setex).toHaveBeenCalledWith(
|
||||
"tts_key_rotation",
|
||||
604800,
|
||||
"1",
|
||||
);
|
||||
});
|
||||
|
||||
test("rotates after an API rejection and logs the failure", async () => {
|
||||
const deps = dependencies();
|
||||
deps.generateContent
|
||||
.mockRejectedValueOnce(new Error("quota"))
|
||||
.mockResolvedValueOnce("audio");
|
||||
const response = await createTtsHandler(deps)(request());
|
||||
expect(response.status).toBe(200);
|
||||
expect(
|
||||
deps.generateContent.mock.calls.map((call) => call[0].apiKey),
|
||||
).toEqual(["first", "second"]);
|
||||
expect(deps.logError).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
test("continues rotating when an API response has no audio", async () => {
|
||||
const deps = dependencies();
|
||||
deps.generateContent
|
||||
.mockResolvedValueOnce(undefined)
|
||||
.mockResolvedValueOnce("audio");
|
||||
await createTtsHandler(deps)(request());
|
||||
expect(
|
||||
deps.generateContent.mock.calls.map((call) => call[0].apiKey),
|
||||
).toEqual(["first", "second"]);
|
||||
});
|
||||
|
||||
test("returns 500 after ten empty responses and does not cache", async () => {
|
||||
const deps = dependencies();
|
||||
deps.generateContent.mockResolvedValue(undefined);
|
||||
const response = await createTtsHandler(deps)(request());
|
||||
expect(response.status).toBe(500);
|
||||
expect(await response.json()).toEqual({ error: "Response is empty" });
|
||||
expect(deps.generateContent).toHaveBeenCalledTimes(10);
|
||||
expect(deps.redis.setex).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
test.each([
|
||||
"not-a-number",
|
||||
"-2",
|
||||
"99",
|
||||
])("normalizes invalid or out-of-range rotation %p", async (rotation) => {
|
||||
const deps = dependencies();
|
||||
deps.redis.get.mockResolvedValueOnce(null).mockResolvedValueOnce(rotation);
|
||||
await createTtsHandler(deps)(request());
|
||||
const expected = rotation === "99" ? "second" : "first";
|
||||
expect(deps.generateContent.mock.calls[0]?.[0].apiKey).toBe(expected);
|
||||
});
|
||||
|
||||
test("defaults rotation to zero when reading it rejects", async () => {
|
||||
const deps = dependencies();
|
||||
deps.redis.get
|
||||
.mockResolvedValueOnce(null)
|
||||
.mockRejectedValueOnce(new Error("redis unavailable"));
|
||||
await createTtsHandler(deps)(request());
|
||||
expect(deps.generateContent.mock.calls[0]?.[0].apiKey).toBe("first");
|
||||
});
|
||||
|
||||
test("wraps PCM bytes in a mono 24kHz WAV stream", async () => {
|
||||
const pcm = new Uint8Array([1, 2, 3, 4]);
|
||||
const wav = new Uint8Array(
|
||||
await new Response(
|
||||
pcmToWav(Buffer.from(pcm).toString("base64")),
|
||||
).arrayBuffer(),
|
||||
);
|
||||
|
||||
expect(new TextDecoder().decode(wav.slice(0, 4))).toBe("RIFF");
|
||||
expect(new TextDecoder().decode(wav.slice(8, 12))).toBe("WAVE");
|
||||
expect(wav.slice(-pcm.length)).toEqual(pcm);
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,134 @@
|
||||
import { PassThrough, Readable } from "node:stream";
|
||||
import { hash as bunHash } from "bun";
|
||||
import wav from "wav";
|
||||
import z from "zod";
|
||||
|
||||
type TtsRedis = {
|
||||
get(key: string): Promise<string | null>;
|
||||
setex(key: string, ttl: number, value: string): Promise<unknown>;
|
||||
};
|
||||
|
||||
type GenerateContent = (options: {
|
||||
apiKey: string;
|
||||
model: string;
|
||||
message: string;
|
||||
voice: string;
|
||||
}) => Promise<string | undefined>;
|
||||
|
||||
type TtsDependencies = {
|
||||
redis: TtsRedis;
|
||||
widgetKey?: string;
|
||||
apiKeys?: string;
|
||||
generateContent: GenerateContent;
|
||||
hash?: (message: string) => string;
|
||||
random?: () => number;
|
||||
queue?: (callback: () => void) => void;
|
||||
toWav?: (audioData: string) => BodyInit;
|
||||
log?: (...values: unknown[]) => void;
|
||||
logError?: (...values: unknown[]) => void;
|
||||
};
|
||||
|
||||
const MODELS = ["gemini-2.5-flash-preview-tts"] as const;
|
||||
const MAX_TRIES = 10;
|
||||
const voices = ["Zephyr", "Charon", "Fenrir", "Kore", "Sulafat"];
|
||||
const messageSchema = z.string().max(1000);
|
||||
|
||||
export function createTtsHandler({
|
||||
redis,
|
||||
widgetKey,
|
||||
apiKeys = "",
|
||||
generateContent,
|
||||
hash = (message) => bunHash(message).toString(36),
|
||||
random = Math.random,
|
||||
queue = queueMicrotask,
|
||||
toWav = pcmToWav,
|
||||
log = console.log,
|
||||
logError = console.error,
|
||||
}: TtsDependencies) {
|
||||
const choices = MODELS.flatMap((model) =>
|
||||
apiKeys
|
||||
.split(",")
|
||||
.map((key) => key.trim())
|
||||
.filter(Boolean)
|
||||
.map((apiKey) => ({ apiKey, model })),
|
||||
);
|
||||
|
||||
return async function GET(request: Request) {
|
||||
if (!widgetKey || choices.length === 0)
|
||||
return Response.json({ error: "TTS is not configured" }, { status: 503 });
|
||||
|
||||
const params = new URL(request.url).searchParams;
|
||||
const parsed = messageSchema.safeParse(params.get("message"));
|
||||
if (params.get("key") !== widgetKey || !parsed.success)
|
||||
return new Response(
|
||||
!parsed.success ? z.prettifyError(parsed.error) : "Invalid key",
|
||||
{ status: 400 },
|
||||
);
|
||||
|
||||
const message = parsed.data;
|
||||
const textHash = hash(message);
|
||||
const cached = await redis.get(`tts:${textHash}`);
|
||||
log(
|
||||
`[TTS] ${textHash}:${cached ? "cached" : "fresh"} "${message.slice(0, 200)}${message.length > 200 ? "..." : ""}"`,
|
||||
);
|
||||
if (cached)
|
||||
return new Response(toWav(cached), {
|
||||
headers: { "Content-Type": "audio/wav" },
|
||||
});
|
||||
|
||||
let currentRotation = await redis
|
||||
.get("tts_key_rotation")
|
||||
.then((value) => (value ? Number(value) : 0))
|
||||
.catch(() => 0);
|
||||
if (!Number.isInteger(currentRotation) || currentRotation < 0)
|
||||
currentRotation = 0;
|
||||
currentRotation %= choices.length;
|
||||
|
||||
let audioData: string | undefined;
|
||||
for (let attempt = 0; attempt < MAX_TRIES; attempt++) {
|
||||
const choice = choices[currentRotation];
|
||||
log(
|
||||
`[TTS] Trying key ${choice.apiKey.slice(0, 8)}:${choice.model} (${currentRotation}/${choices.length} #${attempt + 1})...`,
|
||||
);
|
||||
try {
|
||||
audioData = await generateContent({
|
||||
...choice,
|
||||
message,
|
||||
voice: voices[Math.floor(random() * voices.length)],
|
||||
});
|
||||
if (audioData) break;
|
||||
} catch (error) {
|
||||
logError(`TTS API key failed: ${choice.apiKey.slice(0, 8)}...`, error);
|
||||
}
|
||||
currentRotation = (currentRotation + 1) % choices.length;
|
||||
}
|
||||
log("[TTS] Got response?", audioData ? "yes" : "no");
|
||||
|
||||
if (!audioData)
|
||||
return Response.json({ error: "Response is empty" }, { status: 500 });
|
||||
|
||||
queue(() => {
|
||||
void redis.setex(`tts:${textHash}`, 604800, audioData);
|
||||
void redis.setex("tts_key_rotation", 604800, `${currentRotation}`);
|
||||
});
|
||||
return new Response(toWav(audioData), {
|
||||
headers: { "Content-Type": "audio/wav" },
|
||||
});
|
||||
};
|
||||
}
|
||||
|
||||
export function pcmToWav(audioData: string) {
|
||||
const pcmStream = Readable.from([Buffer.from(audioData, "base64")]);
|
||||
const writer = new wav.Writer({
|
||||
channels: 1,
|
||||
sampleRate: 24000,
|
||||
bitDepth: 16,
|
||||
});
|
||||
const out = new PassThrough();
|
||||
|
||||
pcmStream.pipe(writer);
|
||||
writer.on("data", (chunk) => out.write(chunk));
|
||||
writer.on("end", () => out.end());
|
||||
writer.on("error", () => out.end());
|
||||
return Readable.toWeb(out) as unknown as globalThis.ReadableStream;
|
||||
}
|
||||
+18
-114
@@ -1,119 +1,23 @@
|
||||
import { GoogleGenAI } from "@google/genai";
|
||||
import { hash } from "bun";
|
||||
import { PassThrough, Readable } from "stream";
|
||||
import wav from "wav";
|
||||
import z from "zod";
|
||||
import { redis } from "@/lib/db/redis";
|
||||
import { createTtsHandler } from "./handler";
|
||||
|
||||
const { DONATE_WIDGET_KEY, GEMINI_TTS_API_KEY } = process.env as Record<
|
||||
string,
|
||||
string
|
||||
>;
|
||||
|
||||
const MODELS = [
|
||||
"gemini-2.5-flash-preview-tts",
|
||||
// "gemini-3.1-flash-preview-tts",
|
||||
] as const;
|
||||
const API_KEYS = GEMINI_TTS_API_KEY.split(",")
|
||||
.map((k) => k.trim())
|
||||
.filter(Boolean);
|
||||
const POSSIBLE_CHOICES = MODELS.flatMap((m) =>
|
||||
API_KEYS.map((k) => `${k}:${m}` as `${string}:${(typeof MODELS)[number]}`),
|
||||
);
|
||||
const MAX_TRIES = 10;
|
||||
|
||||
const Schema = z.object({
|
||||
message: z.string().max(1000),
|
||||
key: z.literal(DONATE_WIDGET_KEY),
|
||||
});
|
||||
|
||||
const voices = ["Zephyr", "Charon", "Fenrir", "Kore", "Sulafat"];
|
||||
|
||||
export async function GET(request: Request) {
|
||||
const data = Schema.safeParse(
|
||||
Object.fromEntries(new URL(request.url).searchParams.entries()),
|
||||
);
|
||||
if (data.error)
|
||||
return new Response(z.prettifyError(data.error), { status: 400 });
|
||||
const {
|
||||
data: { message },
|
||||
} = data;
|
||||
|
||||
const textHash = hash(message).toString(36);
|
||||
const cached = await redis!.get(`tts:${textHash}`);
|
||||
console.log(
|
||||
`[TTS] ${textHash}:${cached ? "cached" : "fresh"} "${message.slice(0, 200)}${message.length > 200 ? "..." : ""}"`,
|
||||
);
|
||||
if (cached)
|
||||
return new Response(pcmToWav(cached), {
|
||||
headers: { "Content-Type": "audio/wav" },
|
||||
});
|
||||
|
||||
let audioData: string | undefined;
|
||||
let currentRotation = await redis!
|
||||
.get("tts_key_rotation")
|
||||
.then((e) => (e ? Number(e) : 0))
|
||||
.catch(() => 0);
|
||||
for (
|
||||
let i = 0;
|
||||
i < MAX_TRIES;
|
||||
i++, currentRotation = (currentRotation + 1) % POSSIBLE_CHOICES.length
|
||||
) {
|
||||
const [apiKey, model] = POSSIBLE_CHOICES[currentRotation].split(":");
|
||||
console.log(
|
||||
`[TTS] Trying key ${apiKey.slice(0, 8)}:${model} (${currentRotation}/${POSSIBLE_CHOICES.length} #${i + 1})...`,
|
||||
);
|
||||
export const GET = createTtsHandler({
|
||||
redis: redis!,
|
||||
widgetKey: process.env.DONATE_WIDGET_KEY,
|
||||
apiKeys: process.env.GEMINI_TTS_API_KEY,
|
||||
async generateContent({ apiKey, model, message, voice }) {
|
||||
const client = new GoogleGenAI({ apiKey });
|
||||
|
||||
try {
|
||||
const response = await client.models.generateContent({
|
||||
model,
|
||||
contents: [{ parts: [{ text: message }] }],
|
||||
config: {
|
||||
responseModalities: ["AUDIO"],
|
||||
speechConfig: {
|
||||
voiceConfig: {
|
||||
prebuiltVoiceConfig: {
|
||||
voiceName: voices[Math.floor(Math.random() * voices.length)],
|
||||
},
|
||||
},
|
||||
},
|
||||
const response = await client.models.generateContent({
|
||||
model,
|
||||
contents: [{ parts: [{ text: message }] }],
|
||||
config: {
|
||||
responseModalities: ["AUDIO"],
|
||||
speechConfig: {
|
||||
voiceConfig: { prebuiltVoiceConfig: { voiceName: voice } },
|
||||
},
|
||||
});
|
||||
|
||||
audioData = response.candidates?.[0].content?.parts?.[0].inlineData?.data;
|
||||
if (audioData) break;
|
||||
} catch (e) {
|
||||
console.error(`TTS API key failed: ${apiKey.slice(0, 8)}...`, e);
|
||||
}
|
||||
}
|
||||
console.log("[TTS] Got response?", audioData ? "yes" : "no");
|
||||
|
||||
if (!audioData)
|
||||
return Response.json({ error: "Response is empty" }, { status: 500 });
|
||||
queueMicrotask(() => {
|
||||
redis!.setex(`tts:${textHash}`, 604800, audioData);
|
||||
redis!.setex("tts_key_rotation", 604800, `${currentRotation}`);
|
||||
});
|
||||
return new Response(pcmToWav(audioData), {
|
||||
headers: { "Content-Type": "audio/wav" },
|
||||
});
|
||||
}
|
||||
|
||||
function pcmToWav(audioData: string) {
|
||||
const audioBuffer = Buffer.from(audioData, "base64");
|
||||
const pcmStream = Readable.from([audioBuffer]);
|
||||
const writer = new wav.Writer({
|
||||
channels: 1,
|
||||
sampleRate: 24000,
|
||||
bitDepth: 16,
|
||||
});
|
||||
|
||||
const out = new PassThrough();
|
||||
|
||||
pcmStream.pipe(writer);
|
||||
writer.on("data", (chunk) => out.write(chunk));
|
||||
writer.on("end", () => out.end());
|
||||
writer.on("error", () => out.end());
|
||||
return Readable.toWeb(out) as unknown as globalThis.ReadableStream;
|
||||
}
|
||||
},
|
||||
});
|
||||
return response.candidates?.[0].content?.parts?.[0].inlineData?.data;
|
||||
},
|
||||
});
|
||||
|
||||
Reference in New Issue
Block a user