Skip to content

Commit 874783b

Browse files
jessicagamiojessicagamio
committed
test(llmobs/openai-java): migrate session_id propagation test to JUnit 5 + Java
1 parent 4bdfc39 commit 874783b

2 files changed

Lines changed: 130 additions & 63 deletions

File tree

‎dd-java-agent/instrumentation/openai-java/openai-java-3.0/src/test/groovy/SessionIdPropagationTest.groovy‎

Lines changed: 0 additions & 63 deletions
This file was deleted.

‎dd-java-agent/instrumentation/openai-java/openai-java-3.0/src/test/java/datadog/trace/instrumentation/openai_java/SessionIdPropagationForkedTest.java‎

Lines changed: 130 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,130 @@
1+
package datadog.trace.instrumentation.openai_java;
2+
3+
import static org.junit.jupiter.api.Assertions.assertEquals;
4+
import static org.junit.jupiter.api.Assertions.assertNotNull;
5+
import static org.junit.jupiter.api.Assertions.assertNull;
6+
7+
import com.openai.client.OpenAIClient;
8+
import com.openai.client.okhttp.OpenAIOkHttpClient;
9+
import com.openai.credential.BearerTokenCredential;
10+
import com.openai.models.ChatModel;
11+
import com.openai.models.chat.completions.ChatCompletionCreateParams;
12+
import com.sun.net.httpserver.HttpServer;
13+
import datadog.context.ContextScope;
14+
import datadog.trace.agent.test.AbstractInstrumentationTest;
15+
import datadog.trace.api.llmobs.LLMObsContext;
16+
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
17+
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
18+
import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
19+
import datadog.trace.core.DDSpan;
20+
import datadog.trace.junit.utils.config.WithConfig;
21+
import java.io.IOException;
22+
import java.net.InetSocketAddress;
23+
import java.util.List;
24+
import org.junit.jupiter.api.AfterAll;
25+
import org.junit.jupiter.api.BeforeAll;
26+
import org.junit.jupiter.api.Test;
27+
28+
/**
29+
* Verifies that auto-instrumented openai.request spans inherit session_id from an active LLMObs
30+
* parent context. Forked + @WithConfig used together so the LLMObs system property is in place
31+
* before the agent installs and there's no leakage from prior test state.
32+
*
33+
* <p>The mock OpenAI backend returns a minimal 200 response — the test asserts on the span tag set
34+
* by OpenAiDecorator.afterStart(), which runs before the HTTP response is parsed, so the response
35+
* body shape doesn't matter for what's being tested.
36+
*/
37+
@WithConfig(key = "llmobs.enabled", value = "true")
38+
class SessionIdPropagationForkedTest extends AbstractInstrumentationTest {
39+
40+
private static HttpServer mockServer;
41+
private static OpenAIClient openAiClient;
42+
43+
@BeforeAll
44+
static void setupMockOpenAi() throws IOException {
45+
mockServer = HttpServer.create(new InetSocketAddress("localhost", 0), 0);
46+
mockServer.createContext(
47+
"/v1/",
48+
exchange -> {
49+
exchange.sendResponseHeaders(200, -1);
50+
exchange.close();
51+
});
52+
mockServer.start();
53+
54+
openAiClient =
55+
OpenAIOkHttpClient.builder()
56+
.baseUrl(
57+
"http://"
58+
+ mockServer.getAddress().getHostString()
59+
+ ":"
60+
+ mockServer.getAddress().getPort()
61+
+ "/v1")
62+
.credential(BearerTokenCredential.create(""))
63+
.build();
64+
}
65+
66+
@AfterAll
67+
static void tearDownMockOpenAi() {
68+
if (mockServer != null) {
69+
mockServer.stop(0);
70+
mockServer = null;
71+
}
72+
openAiClient = null;
73+
}
74+
75+
@Test
76+
void openAiRequestSpanInheritsSessionIdFromActiveContext() throws Exception {
77+
String expectedSessionId = "session-propagation-test-abc";
78+
79+
AgentSpan parentSpan = AgentTracer.startSpan("test", "parent");
80+
try (AgentScope ignored1 = AgentTracer.activateSpan(parentSpan)) {
81+
try (ContextScope ignored2 =
82+
LLMObsContext.attach(parentSpan.context(), expectedSessionId)) {
83+
try {
84+
openAiClient.chat().completions().create(buildMinimalChatParams());
85+
} catch (Exception ignored) {
86+
// Mock server returns no body — the SDK may throw on parse. The span we care about
87+
// is already created by the instrumentation advice before this point.
88+
}
89+
}
90+
} finally {
91+
parentSpan.finish();
92+
}
93+
94+
writer.waitForTraces(1);
95+
DDSpan openAiSpan = findSpanByOperationName(writer, "openai.request");
96+
assertNotNull(openAiSpan, "openai.request span should have been created");
97+
assertEquals(expectedSessionId, openAiSpan.getTag("_ml_obs_tag.session_id"));
98+
}
99+
100+
@Test
101+
void openAiRequestSpanHasNoSessionIdWhenNoLlmObsContext() throws Exception {
102+
try {
103+
openAiClient.chat().completions().create(buildMinimalChatParams());
104+
} catch (Exception ignored) {
105+
// Mock server returns no body — the SDK may throw on parse. The span we care about
106+
// is already created by the instrumentation advice before this point.
107+
}
108+
109+
writer.waitForTraces(1);
110+
DDSpan openAiSpan = findSpanByOperationName(writer, "openai.request");
111+
assertNotNull(openAiSpan, "openai.request span should have been created");
112+
assertNull(openAiSpan.getTag("_ml_obs_tag.session_id"));
113+
}
114+
115+
private static ChatCompletionCreateParams buildMinimalChatParams() {
116+
return ChatCompletionCreateParams.builder()
117+
.model(ChatModel.GPT_4O_MINI)
118+
.addSystemMessage("")
119+
.addUserMessage("")
120+
.build();
121+
}
122+
123+
private static DDSpan findSpanByOperationName(List<List<DDSpan>> traces, String operationName) {
124+
return traces.stream()
125+
.flatMap(List::stream)
126+
.filter(s -> operationName.equals(s.getOperationName().toString()))
127+
.findFirst()
128+
.orElse(null);
129+
}
130+
}

0 commit comments

Comments
 (0)