Skip to content

Commit ca11487

Browse files
realarkrealark
committed
wip
1 parent 5df8ada commit ca11487

7 files changed

Lines changed: 192 additions & 113 deletions

‎src/test/java/dev/braintrust/TestHarness.java‎

Lines changed: 14 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -30,10 +30,12 @@ public class TestHarness {
3030
private static final VCR vcr;
3131

3232
static {
33-
vcr = new VCR(java.util.Map.of(
34-
"https://api.openai.com/v1", "openai",
35-
"https://api.anthropic.com", "anthropic"
36-
));
33+
vcr =
34+
new VCR(
35+
java.util.Map.of(
36+
"https://api.openai.com/v1", "openai",
37+
"https://api.anthropic.com", "anthropic",
38+
"https://generativelanguage.googleapis.com", "google"));
3739
vcr.start();
3840
Runtime.getRuntime().addShutdownHook(new Thread(vcr::stop));
3941
}
@@ -127,6 +129,14 @@ public String anthropicApiKey() {
127129
return vcr.anthropicApiKey();
128130
}
129131

132+
public String geminiBaseUrl() {
133+
return vcr.getUrlForTargetBase("https://generativelanguage.googleapis.com");
134+
}
135+
136+
public String geminiApiKey() {
137+
return vcr.geminiApiKey();
138+
}
139+
130140
/** flush all pending spans and return all spans which have been exported so far */
131141
public List<SpanData> awaitExportedSpans() {
132142
assertTrue(

‎src/test/java/dev/braintrust/VCR.java‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -131,6 +131,18 @@ public synchronized String anthropicApiKey() {
131131
return System.getenv("ANTHROPIC_API_KEY");
132132
}
133133

134+
/**
135+
* Get the Gemini API key. Returns a test key in replay mode, real key from environment in
136+
* record/off modes.
137+
*/
138+
public synchronized String geminiApiKey() {
139+
assertStarted();
140+
if (mode == VcrMode.REPLAY) {
141+
return "test-api-key";
142+
}
143+
return System.getenv("GEMINI_API_KEY");
144+
}
145+
134146
private void startRecording() {
135147
if (mode == VcrMode.RECORD && !recordingStarted) {
136148
targetUrlToMappingsDir.keySet().forEach(this::startRecording);

‎src/test/java/dev/braintrust/instrumentation/genai/BraintrustGenAITest.java‎

Lines changed: 24 additions & 109 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,8 @@
11
package dev.braintrust.instrumentation.genai;
22

3-
import static com.github.tomakehurst.wiremock.client.WireMock.*;
4-
import static com.github.tomakehurst.wiremock.core.WireMockConfiguration.wireMockConfig;
53
import static org.junit.jupiter.api.Assertions.*;
64

75
import com.fasterxml.jackson.databind.ObjectMapper;
8-
import com.github.tomakehurst.wiremock.junit5.WireMockExtension;
96
import com.google.genai.Client;
107
import com.google.genai.types.GenerateContentConfig;
118
import com.google.genai.types.HttpOptions;
@@ -14,67 +11,31 @@
1411
import lombok.SneakyThrows;
1512
import org.junit.jupiter.api.BeforeEach;
1613
import org.junit.jupiter.api.Test;
17-
import org.junit.jupiter.api.extension.RegisterExtension;
1814

1915
public class BraintrustGenAITest {
2016
private static final ObjectMapper JSON_MAPPER = new ObjectMapper();
2117

22-
@RegisterExtension
23-
static WireMockExtension wireMock =
24-
WireMockExtension.newInstance().options(wireMockConfig().dynamicPort()).build();
25-
2618
private TestHarness testHarness;
2719

2820
@BeforeEach
2921
void beforeEach() {
3022
testHarness = TestHarness.setup();
31-
wireMock.resetAll();
3223
}
3324

3425
@Test
3526
@SneakyThrows
3627
void testWrapGemini() {
37-
// Mock the Gemini API response
38-
wireMock.stubFor(
39-
post(urlPathMatching("/v1beta/models/.*:generateContent"))
40-
.willReturn(
41-
aResponse()
42-
.withStatus(200)
43-
.withHeader("Content-Type", "application/json")
44-
.withBody(
45-
"""
46-
{
47-
"candidates": [
48-
{
49-
"content": {
50-
"parts": [
51-
{
52-
"text": "The capital of France is Paris."
53-
}
54-
],
55-
"role": "model"
56-
},
57-
"finishReason": "STOP"
58-
}
59-
],
60-
"usageMetadata": {
61-
"promptTokenCount": 10,
62-
"candidatesTokenCount": 8,
63-
"totalTokenCount": 18
64-
},
65-
"modelVersion": "gemini-2.0-flash-lite"
66-
}
67-
""")));
68-
69-
// Create Gemini client pointing to WireMock server
28+
// Create Gemini client using VCR
7029
HttpOptions httpOptions =
71-
HttpOptions.builder().baseUrl("http://localhost:" + wireMock.getPort()).build();
30+
HttpOptions.builder().baseUrl(testHarness.geminiBaseUrl()).build();
7231

7332
// Wrap with Braintrust instrumentation
7433
var geminiClient =
7534
BraintrustGenAI.wrap(
7635
testHarness.openTelemetry(),
77-
new Client.Builder().apiKey("test-api-key").httpOptions(httpOptions));
36+
new Client.Builder()
37+
.apiKey(testHarness.geminiApiKey())
38+
.httpOptions(httpOptions));
7839

7940
var config = GenerateContentConfig.builder().temperature(0.0f).maxOutputTokens(50).build();
8041

@@ -84,8 +45,7 @@ void testWrapGemini() {
8445

8546
// Verify the response
8647
assertNotNull(response);
87-
wireMock.verify(1, postRequestedFor(urlPathMatching("/v1beta/models/.*:generateContent")));
88-
assertEquals("The capital of France is Paris.", response.text());
48+
assertNotNull(response.text());
8949

9050
// Verify spans were exported
9151
var spans = testHarness.awaitExportedSpans();
@@ -109,9 +69,9 @@ void testWrapGemini() {
10969
String metricsJson = span.getAttributes().get(AttributeKey.stringKey("braintrust.metrics"));
11070
assertNotNull(metricsJson, "braintrust.metrics should be set");
11171
var metrics = JSON_MAPPER.readTree(metricsJson);
112-
assertEquals(10, metrics.get("prompt_tokens").asInt());
113-
assertEquals(8, metrics.get("completion_tokens").asInt());
114-
assertEquals(18, metrics.get("tokens").asInt());
72+
assertTrue(metrics.get("prompt_tokens").asInt() > 0, "prompt_tokens should be > 0");
73+
assertTrue(metrics.get("completion_tokens").asInt() > 0, "completion_tokens should be > 0");
74+
assertTrue(metrics.get("tokens").asInt() > 0, "tokens should be > 0");
11575

11676
// Verify braintrust.span_attributes marks this as an LLM span
11777
String spanAttributesJson =
@@ -135,62 +95,25 @@ void testWrapGemini() {
13595
assertNotNull(outputJson, "braintrust.output_json should be set");
13696
var output = JSON_MAPPER.readTree(outputJson);
13797
assertTrue(output.has("candidates"), "output should have candidates");
138-
assertEquals("STOP", output.get("candidates").get(0).get("finishReason").asText());
139-
assertEquals(
140-
"The capital of France is Paris.",
141-
output.get("candidates")
142-
.get(0)
143-
.get("content")
144-
.get("parts")
145-
.get(0)
146-
.get("text")
147-
.asText());
98+
assertNotNull(output.get("candidates").get(0).get("finishReason"));
99+
assertNotNull(
100+
output.get("candidates").get(0).get("content").get("parts").get(0).get("text"));
148101
}
149102

150103
@Test
151104
@SneakyThrows
152105
void testWrapGeminiAsync() {
153-
// Mock the Gemini API response
154-
wireMock.stubFor(
155-
post(urlPathMatching("/v1beta/models/.*:generateContent"))
156-
.willReturn(
157-
aResponse()
158-
.withStatus(200)
159-
.withHeader("Content-Type", "application/json")
160-
.withBody(
161-
"""
162-
{
163-
"candidates": [
164-
{
165-
"content": {
166-
"parts": [
167-
{
168-
"text": "The capital of Germany is Berlin."
169-
}
170-
],
171-
"role": "model"
172-
},
173-
"finishReason": "STOP"
174-
}
175-
],
176-
"usageMetadata": {
177-
"promptTokenCount": 10,
178-
"candidatesTokenCount": 8,
179-
"totalTokenCount": 18
180-
},
181-
"modelVersion": "gemini-2.0-flash-lite"
182-
}
183-
""")));
184-
185-
// Create Gemini client pointing to WireMock server
106+
// Create Gemini client using VCR
186107
HttpOptions httpOptions =
187-
HttpOptions.builder().baseUrl("http://localhost:" + wireMock.getPort()).build();
108+
HttpOptions.builder().baseUrl(testHarness.geminiBaseUrl()).build();
188109

189110
// Wrap with Braintrust instrumentation
190111
var geminiClient =
191112
BraintrustGenAI.wrap(
192113
testHarness.openTelemetry(),
193-
new Client.Builder().apiKey("test-api-key").httpOptions(httpOptions));
114+
new Client.Builder()
115+
.apiKey(testHarness.geminiApiKey())
116+
.httpOptions(httpOptions));
194117

195118
var config = GenerateContentConfig.builder().temperature(0.0f).maxOutputTokens(50).build();
196119

@@ -203,8 +126,7 @@ void testWrapGeminiAsync() {
203126

204127
// Verify the response
205128
assertNotNull(response);
206-
wireMock.verify(1, postRequestedFor(urlPathMatching("/v1beta/models/.*:generateContent")));
207-
assertEquals("The capital of Germany is Berlin.", response.text());
129+
assertNotNull(response.text());
208130

209131
// Verify spans were exported
210132
var spans = testHarness.awaitExportedSpans();
@@ -228,9 +150,9 @@ void testWrapGeminiAsync() {
228150
String metricsJson = span.getAttributes().get(AttributeKey.stringKey("braintrust.metrics"));
229151
assertNotNull(metricsJson, "braintrust.metrics should be set");
230152
var metrics = JSON_MAPPER.readTree(metricsJson);
231-
assertEquals(10, metrics.get("prompt_tokens").asInt());
232-
assertEquals(8, metrics.get("completion_tokens").asInt());
233-
assertEquals(18, metrics.get("tokens").asInt());
153+
assertTrue(metrics.get("prompt_tokens").asInt() > 0, "prompt_tokens should be > 0");
154+
assertTrue(metrics.get("completion_tokens").asInt() > 0, "completion_tokens should be > 0");
155+
assertTrue(metrics.get("tokens").asInt() > 0, "tokens should be > 0");
234156

235157
// Verify braintrust.span_attributes marks this as an LLM span
236158
String spanAttributesJson =
@@ -254,15 +176,8 @@ void testWrapGeminiAsync() {
254176
assertNotNull(outputJson, "braintrust.output_json should be set");
255177
var output = JSON_MAPPER.readTree(outputJson);
256178
assertTrue(output.has("candidates"), "output should have candidates");
257-
assertEquals("STOP", output.get("candidates").get(0).get("finishReason").asText());
258-
assertEquals(
259-
"The capital of Germany is Berlin.",
260-
output.get("candidates")
261-
.get(0)
262-
.get("content")
263-
.get("parts")
264-
.get(0)
265-
.get("text")
266-
.asText());
179+
assertNotNull(output.get("candidates").get(0).get("finishReason"));
180+
assertNotNull(
181+
output.get("candidates").get(0).get("content").get("parts").get(0).get("text"));
267182
}
268183
}

‎src/test/resources/cassettes/google/__files/v1beta_models_gemini-20-flash-litegeneratecontent-145d189b-0a3e-46a3-a2b4-e517445320f6.json‎

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,35 @@
1+
{
2+
"candidates": [
3+
{
4+
"content": {
5+
"parts": [
6+
{
7+
"text": "The capital of France is **Paris**.\n"
8+
}
9+
],
10+
"role": "model"
11+
},
12+
"finishReason": "STOP",
13+
"avgLogprobs": -0.0092680321799384225
14+
}
15+
],
16+
"usageMetadata": {
17+
"promptTokenCount": 7,
18+
"candidatesTokenCount": 9,
19+
"totalTokenCount": 16,
20+
"promptTokensDetails": [
21+
{
22+
"modality": "TEXT",
23+
"tokenCount": 7
24+
}
25+
],
26+
"candidatesTokensDetails": [
27+
{
28+
"modality": "TEXT",
29+
"tokenCount": 9
30+
}
31+
]
32+
},
33+
"modelVersion": "gemini-2.0-flash-lite",
34+
"responseId": "8LVUaYeADJXQz7IP26vjsQg"
35+
}

‎src/test/resources/cassettes/google/__files/v1beta_models_gemini-20-flash-litegeneratecontent-c4b29720-f406-44ea-a2b1-1f61a05609bb.json‎

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,35 @@
1+
{
2+
"candidates": [
3+
{
4+
"content": {
5+
"parts": [
6+
{
7+
"text": "The capital of Germany is Berlin.\n"
8+
}
9+
],
10+
"role": "model"
11+
},
12+
"finishReason": "STOP",
13+
"avgLogprobs": -0.081099182367324829
14+
}
15+
],
16+
"usageMetadata": {
17+
"promptTokenCount": 7,
18+
"candidatesTokenCount": 8,
19+
"totalTokenCount": 15,
20+
"promptTokensDetails": [
21+
{
22+
"modality": "TEXT",
23+
"tokenCount": 7
24+
}
25+
],
26+
"candidatesTokensDetails": [
27+
{
28+
"modality": "TEXT",
29+
"tokenCount": 8
30+
}
31+
]
32+
},
33+
"modelVersion": "gemini-2.0-flash-lite",
34+
"responseId": "77VUabKfIoGhz7IPm-WN-Qg"
35+
}

‎src/test/resources/cassettes/google/mappings/v1beta_models_gemini-20-flash-litegeneratecontent-145d189b-0a3e-46a3-a2b4-e517445320f6.json‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,36 @@
1+
{
2+
"id" : "145d189b-0a3e-46a3-a2b4-e517445320f6",
3+
"name" : "v1beta_models_gemini-20-flash-litegeneratecontent",
4+
"request" : {
5+
"url" : "/v1beta/models/gemini-2.0-flash-lite:generateContent",
6+
"method" : "POST",
7+
"headers" : {
8+
"Content-Type" : {
9+
"equalTo" : "application/json; charset=UTF-8"
10+
}
11+
},
12+
"bodyPatterns" : [ {
13+
"equalToJson" : "{\"contents\":[{\"parts\":[{\"text\":\"What is the capital of France?\"}],\"role\":\"user\"}],\"generationConfig\":{\"temperature\":0.0,\"maxOutputTokens\":50}}",
14+
"ignoreArrayOrder" : true,
15+
"ignoreExtraElements" : true
16+
} ]
17+
},
18+
"response" : {
19+
"status" : 200,
20+
"bodyFileName" : "v1beta_models_gemini-20-flash-litegeneratecontent-145d189b-0a3e-46a3-a2b4-e517445320f6.json",
21+
"headers" : {
22+
"Content-Type" : "application/json; charset=UTF-8",
23+
"Vary" : [ "Origin", "X-Origin", "Referer" ],
24+
"Date" : "Wed, 31 Dec 2025 05:34:40 GMT",
25+
"Server" : "scaffolding on HTTPServer2",
26+
"X-XSS-Protection" : "0",
27+
"X-Frame-Options" : "SAMEORIGIN",
28+
"X-Content-Type-Options" : "nosniff",
29+
"Server-Timing" : "gfet4t7; dur=442",
30+
"Alt-Svc" : "h3=\":443\"; ma=2592000,h3-29=\":443\"; ma=2592000"
31+
}
32+
},
33+
"uuid" : "145d189b-0a3e-46a3-a2b4-e517445320f6",
34+
"persistent" : true,
35+
"insertionIndex" : 2
36+
}

0 commit comments

Comments
 (0)