From 543b68f785f98f62976655d5f4510bccfd0ffb42 Mon Sep 17 00:00:00 2001 From: cte Date: Mon, 24 Feb 2025 16:09:31 -0800 Subject: [PATCH] Fix tests --- src/api/providers/__tests__/anthropic.test.ts | 76 +++++-------------- .../__tests__/bedrock-converse-format.test.ts | 8 +- .../transform/__tests__/openai-format.test.ts | 6 ++ .../__tests__/vscode-lm-format.test.ts | 3 + 4 files changed, 33 insertions(+), 60 deletions(-) diff --git a/src/api/providers/__tests__/anthropic.test.ts b/src/api/providers/__tests__/anthropic.test.ts index df0050ab9c..ff7bdb4054 100644 --- a/src/api/providers/__tests__/anthropic.test.ts +++ b/src/api/providers/__tests__/anthropic.test.ts @@ -1,50 +1,13 @@ +// npx jest src/api/providers/__tests__/anthropic.test.ts + import { AnthropicHandler } from "../anthropic" import { ApiHandlerOptions } from "../../../shared/api" -import { ApiStream } from "../../transform/stream" -import { Anthropic } from "@anthropic-ai/sdk" -// Mock Anthropic client -const mockBetaCreate = jest.fn() const mockCreate = jest.fn() + jest.mock("@anthropic-ai/sdk", () => { return { Anthropic: jest.fn().mockImplementation(() => ({ - beta: { - promptCaching: { - messages: { - create: mockBetaCreate.mockImplementation(async () => ({ - async *[Symbol.asyncIterator]() { - yield { - type: "message_start", - message: { - usage: { - input_tokens: 100, - output_tokens: 50, - cache_creation_input_tokens: 20, - cache_read_input_tokens: 10, - }, - }, - } - yield { - type: "content_block_start", - index: 0, - content_block: { - type: "text", - text: "Hello", - }, - } - yield { - type: "content_block_delta", - delta: { - type: "text_delta", - text: " world", - }, - } - }, - })), - }, - }, - }, messages: { create: mockCreate.mockImplementation(async (options) => { if (!options.stream) { @@ -65,16 +28,26 @@ jest.mock("@anthropic-ai/sdk", () => { type: "message_start", message: { usage: { - input_tokens: 10, - output_tokens: 5, + input_tokens: 100, + output_tokens: 50, + cache_creation_input_tokens: 20, + cache_read_input_tokens: 10, }, }, } yield { type: "content_block_start", + index: 0, content_block: { type: "text", - text: "Test response", + text: "Hello", + }, + } + yield { + type: "content_block_delta", + delta: { + type: "text_delta", + text: " world", }, } }, @@ -95,7 +68,6 @@ describe("AnthropicHandler", () => { apiModelId: "claude-3-5-sonnet-20241022", } handler = new AnthropicHandler(mockOptions) - mockBetaCreate.mockClear() mockCreate.mockClear() }) @@ -126,17 +98,6 @@ describe("AnthropicHandler", () => { describe("createMessage", () => { const systemPrompt = "You are a helpful assistant." - const messages: Anthropic.Messages.MessageParam[] = [ - { - role: "user", - content: [ - { - type: "text" as const, - text: "Hello!", - }, - ], - }, - ] it("should handle prompt caching for supported models", async () => { const stream = handler.createMessage(systemPrompt, [ @@ -173,9 +134,8 @@ describe("AnthropicHandler", () => { expect(textChunks[0].text).toBe("Hello") expect(textChunks[1].text).toBe(" world") - // Verify beta API was used - expect(mockBetaCreate).toHaveBeenCalled() - expect(mockCreate).not.toHaveBeenCalled() + // Verify API + expect(mockCreate).toHaveBeenCalled() }) }) diff --git a/src/api/transform/__tests__/bedrock-converse-format.test.ts b/src/api/transform/__tests__/bedrock-converse-format.test.ts index c46eb94a2e..fdd29c75bf 100644 --- a/src/api/transform/__tests__/bedrock-converse-format.test.ts +++ b/src/api/transform/__tests__/bedrock-converse-format.test.ts @@ -1,3 +1,5 @@ +// npx jest src/api/transform/__tests__/bedrock-converse-format.test.ts + import { convertToBedrockConverseMessages, convertToAnthropicMessage } from "../bedrock-converse-format" import { Anthropic } from "@anthropic-ai/sdk" import { ContentBlock, ToolResultContentBlock } from "@aws-sdk/client-bedrock-runtime" @@ -187,6 +189,8 @@ describe("bedrock-converse-format", () => { usage: { input_tokens: 10, output_tokens: 20, + cache_creation_input_tokens: null, + cache_read_input_tokens: null, }, }) }) @@ -205,7 +209,7 @@ describe("bedrock-converse-format", () => { expect(result).toEqual({ type: "message", role: "assistant", - content: [{ type: "text", text: "Hello" }], + content: [{ type: "text", text: "Hello", citations: null }], model: "test-model", }) }) @@ -224,7 +228,7 @@ describe("bedrock-converse-format", () => { expect(result).toEqual({ type: "message", role: "assistant", - content: [{ type: "text", text: " world" }], + content: [{ type: "text", text: " world", citations: null }], model: "test-model", }) }) diff --git a/src/api/transform/__tests__/openai-format.test.ts b/src/api/transform/__tests__/openai-format.test.ts index f37d369d70..812208acd1 100644 --- a/src/api/transform/__tests__/openai-format.test.ts +++ b/src/api/transform/__tests__/openai-format.test.ts @@ -1,3 +1,5 @@ +// npx jest src/api/transform/__tests__/openai-format.test.ts + import { convertToOpenAiMessages, convertToAnthropicMessage } from "../openai-format" import { Anthropic } from "@anthropic-ai/sdk" import OpenAI from "openai" @@ -172,11 +174,14 @@ describe("OpenAI Format Transformations", () => { expect(anthropicMessage.content[0]).toEqual({ type: "text", text: "Hello there!", + citations: null, }) expect(anthropicMessage.stop_reason).toBe("end_turn") expect(anthropicMessage.usage).toEqual({ input_tokens: 10, output_tokens: 5, + cache_creation_input_tokens: null, + cache_read_input_tokens: null, }) }) @@ -221,6 +226,7 @@ describe("OpenAI Format Transformations", () => { expect(anthropicMessage.content[0]).toEqual({ type: "text", text: "Let me check the weather.", + citations: null, }) expect(anthropicMessage.content[1]).toEqual({ type: "tool_use", diff --git a/src/api/transform/__tests__/vscode-lm-format.test.ts b/src/api/transform/__tests__/vscode-lm-format.test.ts index b27097fd17..eb800e2b7a 100644 --- a/src/api/transform/__tests__/vscode-lm-format.test.ts +++ b/src/api/transform/__tests__/vscode-lm-format.test.ts @@ -1,3 +1,5 @@ +// npx jest src/api/transform/__tests__/vscode-lm-format.test.ts + import { Anthropic } from "@anthropic-ai/sdk" import * as vscode from "vscode" import { convertToVsCodeLmMessages, convertToAnthropicRole, convertToAnthropicMessage } from "../vscode-lm-format" @@ -216,6 +218,7 @@ describe("vscode-lm-format", () => { expect(result.content[0]).toEqual({ type: "text", text: "Hello", + citations: null, }) expect(result.id).toBe("test-uuid") })