Skip to content

Commit 316a9ce

Browse files
colegottdankcolegottdank
authored
Fix registry tests 2 (#4685)
1 parent 631da4e commit 316a9ce

File tree

9 files changed

+2620
-2011
lines changed

9 files changed

+2620
-2011
lines changed

‎worker/test/ai-gateway/pass-through.spec.ts‎

Lines changed: 457 additions & 0 deletions
Large diffs are not rendered by default.

‎worker/test/ai-gateway/registry-anthropic.spec.ts‎

Lines changed: 647 additions & 1097 deletions
Large diffs are not rendered by default.

‎worker/test/ai-gateway/registry-google.spec.ts‎

Lines changed: 162 additions & 76 deletions
Original file line numberDiff line numberDiff line change
@@ -1,101 +1,187 @@
1-
import { SELF } from "cloudflare:test";
2-
import { describe, it, expect, beforeEach, afterEach } from "vitest";
1+
import { describe, it, beforeEach, vi } from "vitest";
32
import "../setup";
4-
import {
5-
setupTestEnvironment,
6-
cleanupTestEnvironment,
7-
mockGoogleEndpoint,
8-
createAIGatewayRequest,
9-
} from "../test-utils";
3+
import { runGatewayTest } from "./test-framework";
4+
import { createOpenAIMockResponse } from "../test-utils";
5+
6+
// Define auth expectations for Google AI Studio
7+
const googleAuthExpectations = {
8+
headers: {
9+
// Google AI Studio uses OpenAI compatibility mode with Authorization header
10+
Authorization: /^Bearer /,
11+
},
12+
};
1013

1114
describe("Google Registry Tests", () => {
1215
beforeEach(() => {
13-
setupTestEnvironment();
14-
});
15-
16-
afterEach(() => {
17-
cleanupTestEnvironment();
16+
// Clear all mocks between tests
17+
vi.clearAllMocks();
1818
});
1919

2020
describe("BYOK Tests - Google Gemini Models", () => {
2121
// Note: Google models only have the 'google-ai-studio' provider, no vertex/bedrock
2222

2323
// Gemini 2.5 Flash Tests
2424
describe("gemini-2.5-flash", () => {
25-
it("should handle google-ai-studio provider", async () => {
26-
mockGoogleEndpoint("gemini-2.5-flash");
27-
28-
const response = await SELF.fetch(
29-
"https://ai-gateway.helicone.ai/v1/chat/completions",
30-
createAIGatewayRequest("gemini-2.5-flash/google-ai-studio")
31-
);
32-
33-
expect(response.status).toBe(200);
34-
});
35-
36-
it("should auto-select google-ai-studio provider when none specified", async () => {
37-
mockGoogleEndpoint("gemini-2.5-flash");
38-
39-
const response = await SELF.fetch(
40-
"https://ai-gateway.helicone.ai/v1/chat/completions",
41-
createAIGatewayRequest("gemini-2.5-flash")
42-
);
43-
44-
expect(response.status).toBe(200);
45-
});
25+
it("should handle google-ai-studio provider", () =>
26+
runGatewayTest({
27+
model: "gemini-2.5-flash/google-ai-studio",
28+
expected: {
29+
providers: [
30+
{
31+
url: "https://generativelanguage.googleapis.com/v1beta/openai/chat/completions",
32+
response: "success",
33+
model: "gemini-2.5-flash",
34+
data: createOpenAIMockResponse("gemini-2.5-flash"),
35+
expects: googleAuthExpectations,
36+
},
37+
],
38+
finalStatus: 200,
39+
},
40+
}));
41+
42+
it("should auto-select google-ai-studio provider when none specified", () =>
43+
runGatewayTest({
44+
model: "gemini-2.5-flash",
45+
expected: {
46+
providers: [
47+
{
48+
url: "https://generativelanguage.googleapis.com/v1beta/openai/chat/completions",
49+
response: "success",
50+
model: "gemini-2.5-flash",
51+
data: createOpenAIMockResponse("gemini-2.5-flash"),
52+
expects: googleAuthExpectations,
53+
},
54+
],
55+
finalStatus: 200,
56+
},
57+
}));
4658
});
4759

4860
// Gemini 2.5 Flash Lite Tests
4961
describe("gemini-2.5-flash-lite", () => {
50-
it("should handle google-ai-studio provider", async () => {
51-
mockGoogleEndpoint("gemini-2.5-flash-lite");
52-
53-
const response = await SELF.fetch(
54-
"https://ai-gateway.helicone.ai/v1/chat/completions",
55-
createAIGatewayRequest("gemini-2.5-flash-lite/google-ai-studio")
56-
);
57-
58-
expect(response.status).toBe(200);
59-
});
60-
61-
it("should auto-select google-ai-studio provider when none specified", async () => {
62-
mockGoogleEndpoint("gemini-2.5-flash-lite");
63-
64-
const response = await SELF.fetch(
65-
"https://ai-gateway.helicone.ai/v1/chat/completions",
66-
createAIGatewayRequest("gemini-2.5-flash-lite")
67-
);
68-
69-
expect(response.status).toBe(200);
70-
});
62+
it("should handle google-ai-studio provider", () =>
63+
runGatewayTest({
64+
model: "gemini-2.5-flash-lite/google-ai-studio",
65+
expected: {
66+
providers: [
67+
{
68+
url: "https://generativelanguage.googleapis.com/v1beta/openai/chat/completions",
69+
response: "success",
70+
model: "gemini-2.5-flash-lite",
71+
data: createOpenAIMockResponse("gemini-2.5-flash-lite"),
72+
expects: googleAuthExpectations,
73+
},
74+
],
75+
finalStatus: 200,
76+
},
77+
}));
78+
79+
it("should auto-select google-ai-studio provider when none specified", () =>
80+
runGatewayTest({
81+
model: "gemini-2.5-flash-lite",
82+
expected: {
83+
providers: [
84+
{
85+
url: "https://generativelanguage.googleapis.com/v1beta/openai/chat/completions",
86+
response: "success",
87+
model: "gemini-2.5-flash-lite",
88+
data: createOpenAIMockResponse("gemini-2.5-flash-lite"),
89+
expects: googleAuthExpectations,
90+
},
91+
],
92+
finalStatus: 200,
93+
},
94+
}));
7195
});
7296

7397
// Gemini 2.5 Pro Tests
7498
describe("gemini-2.5-pro", () => {
75-
it("should handle google-ai-studio provider", async () => {
76-
mockGoogleEndpoint("gemini-2.5-pro");
77-
78-
const response = await SELF.fetch(
79-
"https://ai-gateway.helicone.ai/v1/chat/completions",
80-
createAIGatewayRequest("gemini-2.5-pro/google-ai-studio")
81-
);
82-
83-
expect(response.status).toBe(200);
84-
});
85-
86-
it("should auto-select google-ai-studio provider when none specified", async () => {
87-
mockGoogleEndpoint("gemini-2.5-pro");
88-
89-
const response = await SELF.fetch(
90-
"https://ai-gateway.helicone.ai/v1/chat/completions",
91-
createAIGatewayRequest("gemini-2.5-pro")
92-
);
93-
94-
expect(response.status).toBe(200);
95-
});
99+
it("should handle google-ai-studio provider", () =>
100+
runGatewayTest({
101+
model: "gemini-2.5-pro/google-ai-studio",
102+
expected: {
103+
providers: [
104+
{
105+
url: "https://generativelanguage.googleapis.com/v1beta/openai/chat/completions",
106+
response: "success",
107+
model: "gemini-2.5-pro",
108+
data: createOpenAIMockResponse("gemini-2.5-pro"),
109+
expects: googleAuthExpectations,
110+
},
111+
],
112+
finalStatus: 200,
113+
},
114+
}));
115+
116+
it("should auto-select google-ai-studio provider when none specified", () =>
117+
runGatewayTest({
118+
model: "gemini-2.5-pro",
119+
expected: {
120+
providers: [
121+
{
122+
url: "https://generativelanguage.googleapis.com/v1beta/openai/chat/completions",
123+
response: "success",
124+
model: "gemini-2.5-pro",
125+
data: createOpenAIMockResponse("gemini-2.5-pro"),
126+
expects: googleAuthExpectations,
127+
},
128+
],
129+
finalStatus: 200,
130+
},
131+
}));
96132
});
97133

98134
// Note: Since Google models only have one provider, fallback tests aren't needed
99135
// as there's nothing to fallback to
136+
137+
describe("Error scenarios", () => {
138+
it("should handle Google AI Studio provider failure", () =>
139+
runGatewayTest({
140+
model: "gemini-2.5-flash/google-ai-studio",
141+
expected: {
142+
providers: [
143+
{
144+
url: "https://generativelanguage.googleapis.com/v1beta/openai/chat/completions",
145+
response: "failure",
146+
statusCode: 500,
147+
errorMessage: "Google AI Studio service unavailable",
148+
},
149+
],
150+
finalStatus: 500,
151+
},
152+
}));
153+
154+
it("should handle rate limiting from Google", () =>
155+
runGatewayTest({
156+
model: "gemini-2.5-pro/google-ai-studio",
157+
expected: {
158+
providers: [
159+
{
160+
url: "https://generativelanguage.googleapis.com/v1beta/openai/chat/completions",
161+
response: "failure",
162+
statusCode: 429,
163+
errorMessage: "Rate limit exceeded",
164+
},
165+
],
166+
finalStatus: 429,
167+
},
168+
}));
169+
170+
it("should handle authentication failure", () =>
171+
runGatewayTest({
172+
model: "gemini-2.5-flash/google-ai-studio",
173+
expected: {
174+
providers: [
175+
{
176+
url: "https://generativelanguage.googleapis.com/v1beta/openai/chat/completions",
177+
response: "failure",
178+
statusCode: 401,
179+
errorMessage: "Invalid API key",
180+
},
181+
],
182+
finalStatus: 401,
183+
},
184+
}));
185+
});
100186
});
101187
});

0 commit comments

Comments
 (0)