Improve code quality

mistralai · Jan 1, 2024 · 1b3d954 · 1b3d954
1 parent f7049e5
commit 1b3d954
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 2 deletions.
diff --git 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
@@ -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'
+  };
+}