available-models-gemini-key.test.ts 2.5 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586
  1. import { describe, expect, test, vi } from "vitest";
  2. import type { Provider } from "@/types/provider";
  3. const undiciRequestMock = vi.fn(async () => {
  4. return {
  5. statusCode: 200,
  6. body: {
  7. json: async () => ({
  8. models: [{ name: "models/gemini-1.5-pro", displayName: "Gemini 1.5 Pro" }],
  9. }),
  10. text: async () => "",
  11. },
  12. };
  13. });
  14. vi.mock("undici", () => {
  15. return {
  16. request: undiciRequestMock,
  17. };
  18. });
  19. vi.mock("@/repository/key", () => {
  20. return {
  21. validateApiKeyAndGetUser: vi.fn(async () => ({
  22. user: { id: 1, providerGroup: null, isEnabled: true, expiresAt: null },
  23. key: { providerGroup: null, name: "test-key" },
  24. })),
  25. };
  26. });
  27. vi.mock("@/lib/proxy-agent", () => {
  28. return {
  29. createProxyAgentForProvider: vi.fn(() => null),
  30. };
  31. });
  32. describe("handleAvailableModels - Gemini key 传参", () => {
  33. test("Gemini 上游请求不应在 URL query 携带 key,应使用 x-goog-api-key 头", async () => {
  34. const { ProxyProviderResolver } = await import("@/app/v1/_lib/proxy/provider-selector");
  35. const { handleAvailableModels } = await import("@/app/v1/_lib/models/available-models");
  36. const geminiProvider = {
  37. id: 1,
  38. name: "gemini",
  39. providerType: "gemini",
  40. url: "https://generativelanguage.googleapis.com/v1beta",
  41. key: "upstream-api-key",
  42. preserveClientIp: false,
  43. allowedModels: null,
  44. } as unknown as Provider;
  45. vi.spyOn(ProxyProviderResolver, "selectProviderByType").mockImplementation(
  46. async (_authState, providerType) => {
  47. if (providerType === "gemini") {
  48. return { provider: geminiProvider, context: {} as any };
  49. }
  50. return { provider: null, context: {} as any };
  51. }
  52. );
  53. const c = {
  54. req: {
  55. path: "/v1/models",
  56. header: (name: string) => {
  57. if (name.toLowerCase() === "x-goog-api-key") return "user-api-key";
  58. return undefined;
  59. },
  60. query: () => undefined,
  61. },
  62. json: (body: unknown, status?: number) => {
  63. return new Response(JSON.stringify(body), {
  64. status: status ?? 200,
  65. headers: { "content-type": "application/json" },
  66. });
  67. },
  68. } as any;
  69. await handleAvailableModels(c);
  70. expect(undiciRequestMock).toHaveBeenCalledTimes(1);
  71. const [url, options] = undiciRequestMock.mock.calls[0] ?? [];
  72. expect(String(url)).not.toContain("key=");
  73. expect(options?.headers?.["x-goog-api-key"]).toBe("upstream-api-key");
  74. });
  75. });