Skip to content

Commit 827e53c

Browse files
koxudaxikoxudaxi
committed
test: Implement 100% test coverage for OpenAI OpenTelemetry support in TypeScript
1 parent 7a3a4fd commit 827e53c

File tree

1 file changed

+347
-0
lines changed

1 file changed

+347
-0
lines changed

‎sdks/typescript/src/utils/stream-wrapper.test.ts‎

Lines changed: 347 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,347 @@
1+
import { describe, it, expect, vi, beforeEach } from 'vitest';
2+
import { StreamWrapper } from './stream-wrapper';
3+
import { GEN_AI_ATTRIBUTES } from './attributes';
4+
5+
describe('StreamWrapper', () => {
6+
let mockSpan: any;
7+
let mockStream: any;
8+
9+
beforeEach(() => {
10+
vi.clearAllMocks();
11+
mockSpan = {
12+
addEvent: vi.fn(),
13+
setAttribute: vi.fn(),
14+
setAttributes: vi.fn(),
15+
end: vi.fn(),
16+
};
17+
});
18+
19+
function createMockStream(chunks: any[]) {
20+
return {
21+
async *[Symbol.asyncIterator]() {
22+
for (const chunk of chunks) {
23+
yield chunk;
24+
}
25+
},
26+
};
27+
}
28+
29+
it('should wrap and iterate through stream chunks', async () => {
30+
const chunks = [
31+
{
32+
choices: [{ delta: { content: 'Hello' } }],
33+
},
34+
{
35+
choices: [{ delta: { content: ' world' } }],
36+
},
37+
{
38+
choices: [{ delta: { content: '!' }, finish_reason: 'stop' }],
39+
},
40+
];
41+
42+
mockStream = createMockStream(chunks);
43+
const wrapper = new StreamWrapper(mockSpan, mockStream);
44+
45+
const receivedChunks = [];
46+
for await (const chunk of wrapper) {
47+
receivedChunks.push(chunk);
48+
}
49+
50+
expect(receivedChunks).toEqual(chunks);
51+
expect(mockSpan.addEvent).toHaveBeenCalledWith('gen_ai.choice', {
52+
'gen_ai.system': 'openai',
53+
index: 0,
54+
finish_reason: 'stop',
55+
message: JSON.stringify({
56+
role: 'assistant',
57+
content: 'Hello world!',
58+
}),
59+
});
60+
expect(mockSpan.setAttribute).toHaveBeenCalledWith(
61+
GEN_AI_ATTRIBUTES.GEN_AI_RESPONSE_FINISH_REASONS,
62+
['stop']
63+
);
64+
expect(mockSpan.end).toHaveBeenCalled();
65+
});
66+
67+
it('should handle stream with usage data', async () => {
68+
const chunks = [
69+
{
70+
choices: [{ delta: { content: 'Test' } }],
71+
},
72+
{
73+
usage: {
74+
prompt_tokens: 10,
75+
completion_tokens: 5,
76+
total_tokens: 15,
77+
},
78+
},
79+
];
80+
81+
mockStream = createMockStream(chunks);
82+
const wrapper = new StreamWrapper(mockSpan, mockStream);
83+
84+
const receivedChunks = [];
85+
for await (const chunk of wrapper) {
86+
receivedChunks.push(chunk);
87+
}
88+
89+
expect(mockSpan.setAttributes).toHaveBeenCalledWith({
90+
[GEN_AI_ATTRIBUTES.GEN_AI_USAGE_INPUT_TOKENS]: 10,
91+
[GEN_AI_ATTRIBUTES.GEN_AI_USAGE_OUTPUT_TOKENS]: 5,
92+
'gen_ai.usage.total_tokens': 15,
93+
});
94+
});
95+
96+
it('should handle empty stream', async () => {
97+
mockStream = createMockStream([]);
98+
const wrapper = new StreamWrapper(mockSpan, mockStream);
99+
100+
const receivedChunks = [];
101+
for await (const chunk of wrapper) {
102+
receivedChunks.push(chunk);
103+
}
104+
105+
expect(receivedChunks).toEqual([]);
106+
expect(mockSpan.addEvent).not.toHaveBeenCalled();
107+
expect(mockSpan.setAttribute).not.toHaveBeenCalled();
108+
expect(mockSpan.setAttributes).not.toHaveBeenCalled();
109+
expect(mockSpan.end).toHaveBeenCalled();
110+
});
111+
112+
it('should handle chunks without choices', async () => {
113+
const chunks = [
114+
{ other: 'data' },
115+
{
116+
usage: {
117+
prompt_tokens: 10,
118+
completion_tokens: 5,
119+
total_tokens: 15,
120+
},
121+
},
122+
];
123+
124+
mockStream = createMockStream(chunks);
125+
const wrapper = new StreamWrapper(mockSpan, mockStream);
126+
127+
const receivedChunks = [];
128+
for await (const chunk of wrapper) {
129+
receivedChunks.push(chunk);
130+
}
131+
132+
expect(receivedChunks).toEqual(chunks);
133+
expect(mockSpan.addEvent).not.toHaveBeenCalled();
134+
expect(mockSpan.setAttributes).toHaveBeenCalled();
135+
});
136+
137+
it('should handle chunks with empty choices array', async () => {
138+
const chunks = [
139+
{ choices: [] },
140+
{ choices: [{ delta: { content: 'Test' } }] },
141+
];
142+
143+
mockStream = createMockStream(chunks);
144+
const wrapper = new StreamWrapper(mockSpan, mockStream);
145+
146+
const receivedChunks = [];
147+
for await (const chunk of wrapper) {
148+
receivedChunks.push(chunk);
149+
}
150+
151+
expect(receivedChunks).toEqual(chunks);
152+
expect(mockSpan.addEvent).toHaveBeenCalledOnce();
153+
});
154+
155+
it('should handle chunks with missing delta', async () => {
156+
const chunks = [
157+
{ choices: [{}] },
158+
{ choices: [{ delta: { content: 'Test' } }] },
159+
];
160+
161+
mockStream = createMockStream(chunks);
162+
const wrapper = new StreamWrapper(mockSpan, mockStream);
163+
164+
const receivedChunks = [];
165+
for await (const chunk of wrapper) {
166+
receivedChunks.push(chunk);
167+
}
168+
169+
expect(mockSpan.addEvent).toHaveBeenCalledWith('gen_ai.choice', {
170+
'gen_ai.system': 'openai',
171+
index: 0,
172+
finish_reason: 'unknown',
173+
message: JSON.stringify({
174+
role: 'assistant',
175+
content: 'Test',
176+
}),
177+
});
178+
});
179+
180+
it('should handle finish_reason without content', async () => {
181+
const chunks = [{ choices: [{ finish_reason: 'stop' }] }];
182+
183+
mockStream = createMockStream(chunks);
184+
const wrapper = new StreamWrapper(mockSpan, mockStream);
185+
186+
const receivedChunks = [];
187+
for await (const chunk of wrapper) {
188+
receivedChunks.push(chunk);
189+
}
190+
191+
expect(mockSpan.setAttribute).toHaveBeenCalledWith(
192+
GEN_AI_ATTRIBUTES.GEN_AI_RESPONSE_FINISH_REASONS,
193+
['stop']
194+
);
195+
expect(mockSpan.addEvent).not.toHaveBeenCalled();
196+
});
197+
198+
it('should handle finish_reason with null value', async () => {
199+
const chunks = [
200+
{ choices: [{ delta: { content: 'Test' }, finish_reason: null }] },
201+
];
202+
203+
mockStream = createMockStream(chunks);
204+
const wrapper = new StreamWrapper(mockSpan, mockStream);
205+
206+
const receivedChunks = [];
207+
for await (const chunk of wrapper) {
208+
receivedChunks.push(chunk);
209+
}
210+
211+
expect(mockSpan.addEvent).toHaveBeenCalledWith('gen_ai.choice', {
212+
'gen_ai.system': 'openai',
213+
index: 0,
214+
finish_reason: 'unknown',
215+
message: JSON.stringify({
216+
role: 'assistant',
217+
content: 'Test',
218+
}),
219+
});
220+
expect(mockSpan.setAttribute).not.toHaveBeenCalledWith(
221+
GEN_AI_ATTRIBUTES.GEN_AI_RESPONSE_FINISH_REASONS,
222+
expect.any(Array)
223+
);
224+
});
225+
226+
it('should concatenate content from multiple chunks', async () => {
227+
const chunks = [
228+
{ choices: [{ delta: { content: 'The' } }] },
229+
{ choices: [{ delta: { content: ' quick' } }] },
230+
{ choices: [{ delta: { content: ' brown' } }] },
231+
{ choices: [{ delta: { content: ' fox' }, finish_reason: 'stop' }] },
232+
];
233+
234+
mockStream = createMockStream(chunks);
235+
const wrapper = new StreamWrapper(mockSpan, mockStream);
236+
237+
const receivedChunks = [];
238+
for await (const chunk of wrapper) {
239+
receivedChunks.push(chunk);
240+
}
241+
242+
expect(mockSpan.addEvent).toHaveBeenCalledWith('gen_ai.choice', {
243+
'gen_ai.system': 'openai',
244+
index: 0,
245+
finish_reason: 'stop',
246+
message: JSON.stringify({
247+
role: 'assistant',
248+
content: 'The quick brown fox',
249+
}),
250+
});
251+
});
252+
253+
it('should handle stream errors', async () => {
254+
const error = new Error('Stream error');
255+
mockStream = {
256+
async *[Symbol.asyncIterator]() {
257+
yield { choices: [{ delta: { content: 'Hello' } }] };
258+
throw error;
259+
},
260+
};
261+
262+
const wrapper = new StreamWrapper(mockSpan, mockStream);
263+
264+
const receivedChunks = [];
265+
await expect(async () => {
266+
for await (const chunk of wrapper) {
267+
receivedChunks.push(chunk);
268+
}
269+
}).rejects.toThrow('Stream error');
270+
271+
expect(receivedChunks).toHaveLength(1);
272+
expect(mockSpan.addEvent).toHaveBeenCalled();
273+
expect(mockSpan.end).toHaveBeenCalled();
274+
});
275+
276+
it('should call end even when break is used', async () => {
277+
const chunks = [
278+
{ choices: [{ delta: { content: 'Hello' } }] },
279+
{ choices: [{ delta: { content: ' world' } }] },
280+
{ choices: [{ delta: { content: '!' } }] },
281+
];
282+
283+
mockStream = createMockStream(chunks);
284+
const wrapper = new StreamWrapper(mockSpan, mockStream);
285+
286+
const receivedChunks = [];
287+
for await (const chunk of wrapper) {
288+
receivedChunks.push(chunk);
289+
if (receivedChunks.length === 2) break;
290+
}
291+
292+
expect(receivedChunks).toHaveLength(2);
293+
expect(mockSpan.end).toHaveBeenCalled();
294+
});
295+
296+
it('should handle multiple choices in a single chunk', async () => {
297+
const chunks = [
298+
{
299+
choices: [
300+
{ delta: { content: 'First' } },
301+
{ delta: { content: 'Second' } },
302+
],
303+
},
304+
];
305+
306+
mockStream = createMockStream(chunks);
307+
const wrapper = new StreamWrapper(mockSpan, mockStream);
308+
309+
const receivedChunks = [];
310+
for await (const chunk of wrapper) {
311+
receivedChunks.push(chunk);
312+
}
313+
314+
expect(mockSpan.addEvent).toHaveBeenCalledWith('gen_ai.choice', {
315+
'gen_ai.system': 'openai',
316+
index: 0,
317+
finish_reason: 'unknown',
318+
message: JSON.stringify({
319+
role: 'assistant',
320+
content: 'First',
321+
}),
322+
});
323+
});
324+
325+
it('should handle usage with missing fields', async () => {
326+
const chunks = [
327+
{
328+
usage: {
329+
prompt_tokens: 10,
330+
},
331+
},
332+
];
333+
334+
mockStream = createMockStream(chunks);
335+
const wrapper = new StreamWrapper(mockSpan, mockStream);
336+
337+
for await (const _chunk of wrapper) {
338+
// Consume the stream
339+
}
340+
341+
expect(mockSpan.setAttributes).toHaveBeenCalledWith({
342+
[GEN_AI_ATTRIBUTES.GEN_AI_USAGE_INPUT_TOKENS]: 10,
343+
[GEN_AI_ATTRIBUTES.GEN_AI_USAGE_OUTPUT_TOKENS]: undefined,
344+
'gen_ai.usage.total_tokens': undefined,
345+
});
346+
});
347+
});

0 commit comments

Comments
 (0)