mirror of
https://github.com/RooVetGit/Roo-Code.git
synced 2026-09-07 08:26:51 +00:00
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:
parent
bf70923914
commit
29b7c8f9b2
1 changed files with 78 additions and 214 deletions
|
|
@ -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)
|
||||
})
|
||||
})
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue