Skip to content

Commit 8a5563b

Browse files
committed
Add MistralHandler constructor validation and unit tests
1 parent 19c38a3 commit 8a5563b

File tree

2 files changed

+130
-1
lines changed

2 files changed

+130
-1
lines changed
Lines changed: 126 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,126 @@
1+
import { MistralHandler } from "../mistral"
2+
import { ApiHandlerOptions, mistralDefaultModelId } from "../../../shared/api"
3+
import { Anthropic } from "@anthropic-ai/sdk"
4+
import { ApiStreamTextChunk } from "../../transform/stream"
5+
6+
// Mock Mistral client
7+
const mockCreate = jest.fn()
8+
jest.mock("@mistralai/mistralai", () => {
9+
return {
10+
Mistral: jest.fn().mockImplementation(() => ({
11+
chat: {
12+
stream: mockCreate.mockImplementation(async (options) => {
13+
const stream = {
14+
[Symbol.asyncIterator]: async function* () {
15+
yield {
16+
data: {
17+
choices: [
18+
{
19+
delta: { content: "Test response" },
20+
index: 0,
21+
},
22+
],
23+
},
24+
}
25+
},
26+
}
27+
return stream
28+
}),
29+
},
30+
})),
31+
}
32+
})
33+
34+
describe("MistralHandler", () => {
35+
let handler: MistralHandler
36+
let mockOptions: ApiHandlerOptions
37+
38+
beforeEach(() => {
39+
mockOptions = {
40+
apiModelId: "codestral-latest", // Update to match the actual model ID
41+
mistralApiKey: "test-api-key",
42+
includeMaxTokens: true,
43+
modelTemperature: 0,
44+
}
45+
handler = new MistralHandler(mockOptions)
46+
mockCreate.mockClear()
47+
})
48+
49+
describe("constructor", () => {
50+
it("should initialize with provided options", () => {
51+
expect(handler).toBeInstanceOf(MistralHandler)
52+
expect(handler.getModel().id).toBe(mockOptions.apiModelId)
53+
})
54+
55+
it("should throw error if API key is missing", () => {
56+
expect(() => {
57+
new MistralHandler({
58+
...mockOptions,
59+
mistralApiKey: undefined,
60+
})
61+
}).toThrow("Mistral API key is required")
62+
})
63+
64+
it("should use custom base URL if provided", () => {
65+
const customBaseUrl = "https://custom.mistral.ai/v1"
66+
const handlerWithCustomUrl = new MistralHandler({
67+
...mockOptions,
68+
mistralCodestralUrl: customBaseUrl,
69+
})
70+
expect(handlerWithCustomUrl).toBeInstanceOf(MistralHandler)
71+
})
72+
})
73+
74+
describe("getModel", () => {
75+
it("should return correct model info", () => {
76+
const model = handler.getModel()
77+
expect(model.id).toBe(mockOptions.apiModelId)
78+
expect(model.info).toBeDefined()
79+
expect(model.info.supportsPromptCache).toBe(false)
80+
})
81+
})
82+
83+
describe("createMessage", () => {
84+
const systemPrompt = "You are a helpful assistant."
85+
const messages: Anthropic.Messages.MessageParam[] = [
86+
{
87+
role: "user",
88+
content: [{ type: "text", text: "Hello!" }],
89+
},
90+
]
91+
92+
it("should create message successfully", async () => {
93+
const iterator = handler.createMessage(systemPrompt, messages)
94+
const result = await iterator.next()
95+
96+
expect(mockCreate).toHaveBeenCalledWith({
97+
model: mockOptions.apiModelId,
98+
messages: expect.any(Array),
99+
maxTokens: expect.any(Number),
100+
temperature: 0,
101+
})
102+
103+
expect(result.value).toBeDefined()
104+
expect(result.done).toBe(false)
105+
})
106+
107+
it("should handle streaming response correctly", async () => {
108+
const iterator = handler.createMessage(systemPrompt, messages)
109+
const results: ApiStreamTextChunk[] = []
110+
111+
for await (const chunk of iterator) {
112+
if ("text" in chunk) {
113+
results.push(chunk as ApiStreamTextChunk)
114+
}
115+
}
116+
117+
expect(results.length).toBeGreaterThan(0)
118+
expect(results[0].text).toBe("Test response")
119+
})
120+
121+
it("should handle errors gracefully", async () => {
122+
mockCreate.mockRejectedValueOnce(new Error("API Error"))
123+
await expect(handler.createMessage(systemPrompt, messages).next()).rejects.toThrow("API Error")
124+
})
125+
})
126+
})

src/api/providers/mistral.ts

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,9 +21,12 @@ export class MistralHandler implements ApiHandler {
2121
private client: Mistral
2222

2323
constructor(options: ApiHandlerOptions) {
24+
if (!options.mistralApiKey) {
25+
throw new Error("Mistral API key is required")
26+
}
27+
2428
this.options = options
2529
const baseUrl = this.getBaseUrl()
26-
// OR Option 2: Using the built-in debug logger
2730
console.debug(`[Roo Code] MistralHandler using baseUrl: ${baseUrl}`)
2831
this.client = new Mistral({
2932
serverURL: baseUrl,

0 commit comments

Comments
 (0)