Skip to content

Commit 277e296

Browse files
easonliang28easonliang28
committed
fix(tests): fix abortSignal test structure and add taskId metadata (PR Zoo-Code-Org#434)
1 parent c59d61b commit 277e296

5 files changed

Lines changed: 125 additions & 97 deletions

File tree

‎src/api/providers/__tests__/deepseek.spec.ts‎

Lines changed: 33 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -638,41 +638,43 @@ describe("DeepSeekHandler", () => {
638638
expect(toolCallChunks.length).toBeGreaterThan(0)
639639
expect(toolCallChunks[0].name).toBe("get_weather")
640640

641-
it("should pass abortSignal to chat.completions.create when provided in metadata", async () => {
642-
const handler = new DeepSeekHandler({ ...mockOptions, apiKey: "test-key" })
643-
const systemPrompt = "You are a helpful assistant."
644-
const messages: Anthropic.Messages.MessageParam[] = [
645-
{ role: "user", content: [{ type: "text" as const, text: "Hello!" }] },
646-
]
647-
648-
const controller = new AbortController()
649-
const mockAbortSignal = controller.signal
650-
651-
await handler.createMessage(systemPrompt, messages, { abortSignal: mockAbortSignal })
652-
for await (const _chunk of handler.createMessage(systemPrompt, messages)) {
653-
break
654-
}
641+
describe("abortSignal support", () => {
642+
it("should pass abortSignal to chat.completions.create when provided in metadata", async () => {
643+
const handler = new DeepSeekHandler({ ...mockOptions, apiKey: "test-key" })
644+
const systemPrompt = "You are a helpful assistant."
645+
const messages: Anthropic.Messages.MessageParam[] = [
646+
{ role: "user", content: [{ type: "text" as const, text: "Hello!" }] },
647+
]
648+
649+
const controller = new AbortController()
650+
const mockAbortSignal = controller.signal
651+
652+
await handler.createMessage(systemPrompt, messages, { taskId: "test", abortSignal: mockAbortSignal })
653+
for await (const _chunk of handler.createMessage(systemPrompt, messages)) {
654+
break
655+
}
655656

656-
expect(mockCreate).toHaveBeenCalled()
657-
const callArgs = mockCreate.mock.calls[0][0]
658-
expect(callArgs.signal).toBe(mockAbortSignal)
659-
})
657+
expect(mockCreate).toHaveBeenCalled()
658+
const callArgs = mockCreate.mock.calls[0][0]
659+
expect(callArgs.signal).toBe(mockAbortSignal)
660+
})
660661

661-
it("should not include signal when abortSignal is not provided", async () => {
662-
const handler = new DeepSeekHandler({ ...mockOptions, apiKey: "test-key" })
663-
const systemPrompt = "You are a helpful assistant."
664-
const messages: Anthropic.Messages.MessageParam[] = [
665-
{ role: "user", content: [{ type: "text" as const, text: "Hello!" }] },
666-
]
662+
it("should not include signal when abortSignal is not provided", async () => {
663+
const handler = new DeepSeekHandler({ ...mockOptions, apiKey: "test-key" })
664+
const systemPrompt = "You are a helpful assistant."
665+
const messages: Anthropic.Messages.MessageParam[] = [
666+
{ role: "user", content: [{ type: "text" as const, text: "Hello!" }] },
667+
]
667668

668-
await handler.createMessage(systemPrompt, messages)
669-
for await (const _chunk of handler.createMessage(systemPrompt, messages)) {
670-
break
671-
}
669+
await handler.createMessage(systemPrompt, messages)
670+
for await (const _chunk of handler.createMessage(systemPrompt, messages)) {
671+
break
672+
}
672673

673-
expect(mockCreate).toHaveBeenCalled()
674-
const callArgs = mockCreate.mock.calls[0][0]
675-
expect(callArgs.signal).toBeUndefined()
674+
expect(mockCreate).toHaveBeenCalled()
675+
const callArgs = mockCreate.mock.calls[0][0]
676+
expect(callArgs.signal).toBeUndefined()
677+
})
676678
})
677679
})
678680
})

‎src/api/providers/__tests__/openai.spec.ts‎

Lines changed: 33 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -1199,41 +1199,43 @@ describe("OpenAiHandler", () => {
11991199
)
12001200
})
12011201

1202-
it("should pass abortSignal to chat.completions.create when provided in metadata", async () => {
1203-
const handler = new OpenAiHandler(mockOptions)
1204-
const systemPrompt = "You are a helpful assistant."
1205-
const messages: Anthropic.Messages.MessageParam[] = [
1206-
{ role: "user", content: [{ type: "text" as const, text: "Hello!" }] },
1207-
]
1208-
1209-
const controller = new AbortController()
1210-
const mockAbortSignal = controller.signal
1211-
1212-
await handler.createMessage(systemPrompt, messages, { abortSignal: mockAbortSignal })
1213-
for await (const _chunk of handler.createMessage(systemPrompt, messages)) {
1214-
break
1215-
}
1202+
describe("abortSignal support", () => {
1203+
it("should pass abortSignal to chat.completions.create when provided in metadata", async () => {
1204+
const handler = new OpenAiHandler(mockOptions)
1205+
const systemPrompt = "You are a helpful assistant."
1206+
const messages: Anthropic.Messages.MessageParam[] = [
1207+
{ role: "user", content: [{ type: "text" as const, text: "Hello!" }] },
1208+
]
1209+
1210+
const controller = new AbortController()
1211+
const mockAbortSignal = controller.signal
1212+
1213+
await handler.createMessage(systemPrompt, messages, { taskId: "test", abortSignal: mockAbortSignal })
1214+
for await (const _chunk of handler.createMessage(systemPrompt, messages)) {
1215+
break
1216+
}
12161217

1217-
expect(mockCreate).toHaveBeenCalled()
1218-
const callArgs = mockCreate.mock.calls[0][0]
1219-
expect(callArgs.signal).toBe(mockAbortSignal)
1220-
})
1218+
expect(mockCreate).toHaveBeenCalled()
1219+
const callArgs = mockCreate.mock.calls[0][0]
1220+
expect(callArgs.signal).toBe(mockAbortSignal)
1221+
})
12211222

1222-
it("should not include signal when abortSignal is not provided", async () => {
1223-
const handler = new OpenAiHandler(mockOptions)
1224-
const systemPrompt = "You are a helpful assistant."
1225-
const messages: Anthropic.Messages.MessageParam[] = [
1226-
{ role: "user", content: [{ type: "text" as const, text: "Hello!" }] },
1227-
]
1223+
it("should not include signal when abortSignal is not provided", async () => {
1224+
const handler = new OpenAiHandler(mockOptions)
1225+
const systemPrompt = "You are a helpful assistant."
1226+
const messages: Anthropic.Messages.MessageParam[] = [
1227+
{ role: "user", content: [{ type: "text" as const, text: "Hello!" }] },
1228+
]
12281229

1229-
await handler.createMessage(systemPrompt, messages)
1230-
for await (const _chunk of handler.createMessage(systemPrompt, messages)) {
1231-
break
1232-
}
1230+
await handler.createMessage(systemPrompt, messages)
1231+
for await (const _chunk of handler.createMessage(systemPrompt, messages)) {
1232+
break
1233+
}
12331234

1234-
expect(mockCreate).toHaveBeenCalled()
1235-
const callArgs = mockCreate.mock.calls[0][0]
1236-
expect(callArgs.signal).toBeUndefined()
1235+
expect(mockCreate).toHaveBeenCalled()
1236+
const callArgs = mockCreate.mock.calls[0][0]
1237+
expect(callArgs.signal).toBeUndefined()
1238+
})
12371239
})
12381240
})
12391241
})

‎src/api/providers/__tests__/requesty.spec.ts‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -382,7 +382,10 @@ describe("RequestyHandler", () => {
382382
mockCreate.mockRejectedValue(new Error("Unexpected error"))
383383

384384
await expect(handler.completePrompt("test prompt")).rejects.toThrow("Unexpected error")
385+
})
386+
})
385387

388+
describe("abortSignal support", () => {
386389
it("should pass abortSignal to chat.completions.create when provided in metadata", async () => {
387390
const handler = new RequestyHandler(mockOptions)
388391
const systemPrompt = "You are a helpful assistant."
@@ -393,7 +396,7 @@ describe("RequestyHandler", () => {
393396
const controller = new AbortController()
394397
const mockAbortSignal = controller.signal
395398

396-
await handler.createMessage(systemPrompt, messages, { abortSignal: mockAbortSignal })
399+
await handler.createMessage(systemPrompt, messages, { taskId: "test", abortSignal: mockAbortSignal })
397400
for await (const _chunk of handler.createMessage(systemPrompt, messages)) {
398401
break
399402
}
@@ -419,6 +422,7 @@ describe("RequestyHandler", () => {
419422
const callArgs = mockCreate.mock.calls[0][0]
420423
expect(callArgs.signal).toBeUndefined()
421424
})
425+
})
422426
})
423427
})
424428
})

‎src/api/providers/__tests__/sambanova.spec.ts‎

Lines changed: 36 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -138,41 +138,44 @@ describe("SambaNovaHandler", () => {
138138
await messageGenerator.next()
139139

140140
expect(mockCreate).toHaveBeenCalledWith(
141-
expect.objectContaining({
142-
model: modelId,
143-
max_tokens: modelInfo.maxTokens,
144-
temperature: 0.7,
145-
messages: expect.arrayContaining([{ role: "system", content: systemPrompt }]),
146-
stream: true,
147-
stream_options: { include_usage: true },
148-
}),
149-
expect.any(Object),
150-
)
151-
152-
it("createMessage should pass abortSignal to SambaNova client", async () => {
153-
const handlerWithModel = new SambaNovaHandler({
154-
apiModelId: "Meta-Llama-3.3-70B-Instruct" as SambaNovaModelId,
155-
sambaNovaApiKey: "test-sambanova-api-key",
156-
})
157-
158-
mockCreate.mockImplementationOnce(() => {
159-
return {
160-
[Symbol.asyncIterator]: () => ({
161-
async next() {
162-
return { done: true }
163-
},
141+
expect.objectContaining({
142+
model: modelId,
143+
max_tokens: modelInfo.maxTokens,
144+
temperature: 0.7,
145+
messages: expect.arrayContaining([{ role: "system", content: systemPrompt }]),
146+
stream: true,
147+
stream_options: { include_usage: true },
164148
}),
165-
}
149+
expect.any(Object),
150+
)
166151
})
167152

168-
const controller = new AbortController()
169-
const mockAbortSignal = controller.signal
170-
171-
await handlerWithModel.createMessage("system prompt", [], { abortSignal: mockAbortSignal })
172-
173-
expect(mockCreate).toHaveBeenCalled()
174-
const callArgs = mockCreate.mock.calls[0][0]
175-
expect(callArgs.signal).toBe(mockAbortSignal)
176-
})
153+
describe("abortSignal support", () => {
154+
it("createMessage should pass abortSignal to SambaNova client", async () => {
155+
const handlerWithModel = new SambaNovaHandler({
156+
apiModelId: "Meta-Llama-3.3-70B-Instruct" as SambaNovaModelId,
157+
sambaNovaApiKey: "test-sambanova-api-key",
158+
})
159+
160+
mockCreate.mockImplementationOnce(() => {
161+
return {
162+
[Symbol.asyncIterator]: () => ({
163+
async next() {
164+
return { done: true }
165+
},
166+
}),
167+
}
168+
})
169+
170+
const controller = new AbortController()
171+
const mockAbortSignal = controller.signal
172+
173+
await handlerWithModel.createMessage("system prompt", [], { taskId: "test", abortSignal: mockAbortSignal })
174+
175+
expect(mockCreate).toHaveBeenCalled()
176+
const callArgs = mockCreate.mock.calls[0][0]
177+
expect(callArgs.signal).toBe(mockAbortSignal)
178+
})
179+
})
177180
})
178181
})

‎src/api/providers/__tests__/vercel-ai-gateway.spec.ts‎

Lines changed: 18 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -591,6 +591,23 @@ describe("VercelAiGatewayHandler", () => {
591591
}),
592592
)
593593
})
594+
})
595+
596+
describe("abortSignal support", () => {
597+
beforeEach(() => {
598+
mockCreate.mockImplementation(async () => ({
599+
[Symbol.asyncIterator]: async function* () {
600+
yield {
601+
choices: [{ delta: { content: "Test response" } }],
602+
usage: null,
603+
}
604+
yield {
605+
choices: [{ delta: {} }],
606+
usage: { prompt_tokens: 10, completion_tokens: 5 },
607+
}
608+
},
609+
}))
610+
})
594611

595612
it("should pass abortSignal to streamText when provided in metadata", async () => {
596613
const handler = new VercelAiGatewayHandler({
@@ -601,7 +618,7 @@ describe("VercelAiGatewayHandler", () => {
601618
const controller = new AbortController()
602619
const mockAbortSignal = controller.signal
603620

604-
await handler.createMessage("system prompt", [], { abortSignal: mockAbortSignal })
621+
await handler.createMessage("system prompt", [], { taskId: "test", abortSignal: mockAbortSignal })
605622

606623
expect(mockCreate).toHaveBeenCalled()
607624
const callArgs = mockCreate.mock.calls[0][0]

0 commit comments

Comments
 (0)