fix: update claude-code-caching.spec.ts to mock OAuth and streaming client

The caching tests were using outdated mocks for the deprecated runClaudeCode function.
Updated to mock claudeCodeOAuthManager and createStreamingMessage instead,
matching the pattern used in claude-code.spec.ts.
This commit is contained in:
Hannes Rudolph 2025-12-13 21:50:16 -07:00
parent bf70923914
commit 29b7c8f9b2

View file

@ -1,79 +1,57 @@
import type { Anthropic } from "@anthropic-ai/sdk"
import { ClaudeCodeHandler } from "../claude-code"
import { runClaudeCode } from "../../../integrations/claude-code/run"
import type { ApiHandlerOptions } from "../../../shared/api"
import type { ClaudeCodeMessage } from "../../../integrations/claude-code/types"
import type { StreamChunk } from "../../../integrations/claude-code/streaming-client"
import type { ApiStreamUsageChunk } from "../../transform/stream"
// Mock the runClaudeCode function
vi.mock("../../../integrations/claude-code/run", () => ({
runClaudeCode: vi.fn(),
// Mock the OAuth manager
vi.mock("../../../integrations/claude-code/oauth", () => ({
claudeCodeOAuthManager: {
getAccessToken: vi.fn(),
getEmail: vi.fn(),
loadCredentials: vi.fn(),
saveCredentials: vi.fn(),
clearCredentials: vi.fn(),
isAuthenticated: vi.fn(),
},
generateUserId: vi.fn(() => "user_abc123_account_def456_session_ghi789"),
}))
// Mock the streaming client
vi.mock("../../../integrations/claude-code/streaming-client", () => ({
createStreamingMessage: vi.fn(),
}))
const { claudeCodeOAuthManager } = await import("../../../integrations/claude-code/oauth")
const { createStreamingMessage } = await import("../../../integrations/claude-code/streaming-client")
const mockGetAccessToken = vi.mocked(claudeCodeOAuthManager.getAccessToken)
const mockCreateStreamingMessage = vi.mocked(createStreamingMessage)
describe("ClaudeCodeHandler - Caching Support", () => {
let handler: ClaudeCodeHandler
const mockOptions: ApiHandlerOptions = {
apiKey: "test-key",
apiModelId: "claude-3-5-sonnet-20241022",
claudeCodePath: "/test/path",
apiModelId: "claude-sonnet-4-5",
}
beforeEach(() => {
handler = new ClaudeCodeHandler(mockOptions)
vi.clearAllMocks()
mockGetAccessToken.mockResolvedValue("test-access-token")
})
it("should collect cache read tokens from API response", async () => {
const mockStream = async function* (): AsyncGenerator<string | ClaudeCodeMessage> {
// Initial system message
const mockStream = async function* (): AsyncGenerator<StreamChunk> {
yield { type: "text", text: "Hello!" }
yield {
type: "system",
subtype: "init",
session_id: "test-session",
tools: [],
mcp_servers: [],
apiKeySource: "user",
} as ClaudeCodeMessage
// Assistant message with cache tokens
const message: Anthropic.Messages.Message = {
id: "msg_123",
type: "message",
role: "assistant",
model: "claude-3-5-sonnet-20241022",
content: [{ type: "text", text: "Hello!", citations: [] }],
usage: {
input_tokens: 100,
output_tokens: 50,
cache_read_input_tokens: 80, // 80 tokens read from cache
cache_creation_input_tokens: 20, // 20 new tokens cached
},
stop_reason: "end_turn",
stop_sequence: null,
type: "usage",
inputTokens: 100,
outputTokens: 50,
cacheReadTokens: 80,
cacheWriteTokens: 20,
}
yield {
type: "assistant",
message,
session_id: "test-session",
} as ClaudeCodeMessage
// Result with cost
yield {
type: "result",
subtype: "success",
result: "success",
total_cost_usd: 0.001,
is_error: false,
duration_ms: 1000,
duration_api_ms: 900,
num_turns: 1,
session_id: "test-session",
} as ClaudeCodeMessage
}
vi.mocked(runClaudeCode).mockReturnValue(mockStream())
mockCreateStreamingMessage.mockReturnValue(mockStream())
const stream = handler.createMessage("System prompt", [{ role: "user", content: "Hello" }])
@ -92,76 +70,29 @@ describe("ClaudeCodeHandler - Caching Support", () => {
})
it("should accumulate cache tokens across multiple messages", async () => {
const mockStream = async function* (): AsyncGenerator<string | ClaudeCodeMessage> {
// Note: The streaming client handles accumulation internally.
// Each usage chunk represents the accumulated totals for that point in the stream.
// This test verifies that we correctly pass through the accumulated values.
const mockStream = async function* (): AsyncGenerator<StreamChunk> {
yield { type: "text", text: "Part 1" }
yield {
type: "system",
subtype: "init",
session_id: "test-session",
tools: [],
mcp_servers: [],
apiKeySource: "user",
} as ClaudeCodeMessage
// First message chunk
const message1: Anthropic.Messages.Message = {
id: "msg_1",
type: "message",
role: "assistant",
model: "claude-3-5-sonnet-20241022",
content: [{ type: "text", text: "Part 1", citations: [] }],
usage: {
input_tokens: 50,
output_tokens: 25,
cache_read_input_tokens: 40,
cache_creation_input_tokens: 10,
},
stop_reason: null,
stop_sequence: null,
type: "usage",
inputTokens: 50,
outputTokens: 25,
cacheReadTokens: 40,
cacheWriteTokens: 10,
}
yield { type: "text", text: "Part 2" }
yield {
type: "assistant",
message: message1,
session_id: "test-session",
} as ClaudeCodeMessage
// Second message chunk
const message2: Anthropic.Messages.Message = {
id: "msg_2",
type: "message",
role: "assistant",
model: "claude-3-5-sonnet-20241022",
content: [{ type: "text", text: "Part 2", citations: [] }],
usage: {
input_tokens: 50,
output_tokens: 25,
cache_read_input_tokens: 30,
cache_creation_input_tokens: 20,
},
stop_reason: "end_turn",
stop_sequence: null,
type: "usage",
inputTokens: 100, // Accumulated: 50 + 50
outputTokens: 50, // Accumulated: 25 + 25
cacheReadTokens: 70, // Accumulated: 40 + 30
cacheWriteTokens: 30, // Accumulated: 10 + 20
}
yield {
type: "assistant",
message: message2,
session_id: "test-session",
} as ClaudeCodeMessage
yield {
type: "result",
subtype: "success",
result: "success",
total_cost_usd: 0.002,
is_error: false,
duration_ms: 2000,
duration_api_ms: 1800,
num_turns: 1,
session_id: "test-session",
} as ClaudeCodeMessage
}
vi.mocked(runClaudeCode).mockReturnValue(mockStream())
mockCreateStreamingMessage.mockReturnValue(mockStream())
const stream = handler.createMessage("System prompt", [{ role: "user", content: "Hello" }])
@ -170,62 +101,29 @@ describe("ClaudeCodeHandler - Caching Support", () => {
chunks.push(chunk)
}
const usageChunk = chunks.find((c) => c.type === "usage" && "totalCost" in c) as ApiStreamUsageChunk | undefined
expect(usageChunk).toBeDefined()
expect(usageChunk!.inputTokens).toBe(100) // 50 + 50
expect(usageChunk!.outputTokens).toBe(50) // 25 + 25
expect(usageChunk!.cacheReadTokens).toBe(70) // 40 + 30
expect(usageChunk!.cacheWriteTokens).toBe(30) // 10 + 20
// Get the last usage chunk which should have accumulated totals
const usageChunks = chunks.filter((c) => c.type === "usage" && "totalCost" in c) as ApiStreamUsageChunk[]
expect(usageChunks.length).toBe(2)
const lastUsageChunk = usageChunks[usageChunks.length - 1]
expect(lastUsageChunk.inputTokens).toBe(100) // 50 + 50
expect(lastUsageChunk.outputTokens).toBe(50) // 25 + 25
expect(lastUsageChunk.cacheReadTokens).toBe(70) // 40 + 30
expect(lastUsageChunk.cacheWriteTokens).toBe(30) // 10 + 20
})
it("should handle missing cache token fields gracefully", async () => {
const mockStream = async function* (): AsyncGenerator<string | ClaudeCodeMessage> {
const mockStream = async function* (): AsyncGenerator<StreamChunk> {
yield { type: "text", text: "Hello!" }
yield {
type: "system",
subtype: "init",
session_id: "test-session",
tools: [],
mcp_servers: [],
apiKeySource: "user",
} as ClaudeCodeMessage
// Message without cache tokens
const message: Anthropic.Messages.Message = {
id: "msg_123",
type: "message",
role: "assistant",
model: "claude-3-5-sonnet-20241022",
content: [{ type: "text", text: "Hello!", citations: [] }],
usage: {
input_tokens: 100,
output_tokens: 50,
cache_read_input_tokens: null,
cache_creation_input_tokens: null,
},
stop_reason: "end_turn",
stop_sequence: null,
type: "usage",
inputTokens: 100,
outputTokens: 50,
// No cache tokens provided
}
yield {
type: "assistant",
message,
session_id: "test-session",
} as ClaudeCodeMessage
yield {
type: "result",
subtype: "success",
result: "success",
total_cost_usd: 0.001,
is_error: false,
duration_ms: 1000,
duration_api_ms: 900,
num_turns: 1,
session_id: "test-session",
} as ClaudeCodeMessage
}
vi.mocked(runClaudeCode).mockReturnValue(mockStream())
mockCreateStreamingMessage.mockReturnValue(mockStream())
const stream = handler.createMessage("System prompt", [{ role: "user", content: "Hello" }])
@ -238,58 +136,24 @@ describe("ClaudeCodeHandler - Caching Support", () => {
expect(usageChunk).toBeDefined()
expect(usageChunk!.inputTokens).toBe(100)
expect(usageChunk!.outputTokens).toBe(50)
expect(usageChunk!.cacheReadTokens).toBe(0)
expect(usageChunk!.cacheWriteTokens).toBe(0)
expect(usageChunk!.cacheReadTokens).toBeUndefined()
expect(usageChunk!.cacheWriteTokens).toBeUndefined()
})
it("should report zero cost for subscription usage", async () => {
const mockStream = async function* (): AsyncGenerator<string | ClaudeCodeMessage> {
// Subscription usage has apiKeySource: "none"
// Claude Code is always subscription-based, cost should always be 0
const mockStream = async function* (): AsyncGenerator<StreamChunk> {
yield { type: "text", text: "Hello!" }
yield {
type: "system",
subtype: "init",
session_id: "test-session",
tools: [],
mcp_servers: [],
apiKeySource: "none",
} as ClaudeCodeMessage
const message: Anthropic.Messages.Message = {
id: "msg_123",
type: "message",
role: "assistant",
model: "claude-3-5-sonnet-20241022",
content: [{ type: "text", text: "Hello!", citations: [] }],
usage: {
input_tokens: 100,
output_tokens: 50,
cache_read_input_tokens: 80,
cache_creation_input_tokens: 20,
},
stop_reason: "end_turn",
stop_sequence: null,
type: "usage",
inputTokens: 100,
outputTokens: 50,
cacheReadTokens: 80,
cacheWriteTokens: 20,
}
yield {
type: "assistant",
message,
session_id: "test-session",
} as ClaudeCodeMessage
yield {
type: "result",
subtype: "success",
result: "success",
total_cost_usd: 0.001, // This should be ignored for subscription usage
is_error: false,
duration_ms: 1000,
duration_api_ms: 900,
num_turns: 1,
session_id: "test-session",
} as ClaudeCodeMessage
}
vi.mocked(runClaudeCode).mockReturnValue(mockStream())
mockCreateStreamingMessage.mockReturnValue(mockStream())
const stream = handler.createMessage("System prompt", [{ role: "user", content: "Hello" }])
@ -300,6 +164,6 @@ describe("ClaudeCodeHandler - Caching Support", () => {
const usageChunk = chunks.find((c) => c.type === "usage" && "totalCost" in c) as ApiStreamUsageChunk | undefined
expect(usageChunk).toBeDefined()
expect(usageChunk!.totalCost).toBe(0) // Should be 0 for subscription usage
expect(usageChunk!.totalCost).toBe(0) // Should always be 0 for Claude Code (subscription-based)
})
})