diff --git a/tests/client.test.js b/tests/client.test.js index 6032c97..8c3ffb5 100644 --- a/tests/client.test.js +++ b/tests/client.test.js @@ -3,6 +3,7 @@ import { mockListModels, mockFetch, mockChatResponseStreamingPayload, + mockEmbeddingRequest, mockEmbeddingResponsePayload, mockChatResponsePayload, mockFetchStream, @@ -65,7 +66,7 @@ describe('Mistral Client', () => { const mockResponse = mockEmbeddingResponsePayload(); globalThis.fetch = mockFetch(200, mockResponse); - const response = await client.listModels(); + const response = await client.embeddings(mockEmbeddingRequest); expect(response).toEqual(mockResponse); }); }); @@ -76,7 +77,7 @@ describe('Mistral Client', () => { const mockResponse = mockEmbeddingResponsePayload(10); globalThis.fetch = mockFetch(200, mockResponse); - const response = await client.listModels(); + const response = await client.embeddings(mockEmbeddingRequest); expect(response).toEqual(mockResponse); }); }); diff --git a/tests/utils.js b/tests/utils.js index 28bb468..8d80084 100644 --- a/tests/utils.js +++ b/tests/utils.js @@ -243,3 +243,14 @@ export function mockEmbeddingResponsePayload(batchSize = 1) { usage: {prompt_tokens: 90, total_tokens: 90, completion_tokens: 0}, }; } + +/** + * Mock embeddings request payload + * @return {Object} + */ +export function mockEmbeddingRequest() { + return { + model: 'mistral-embed', + input: 'embed', + }; +}