From 1b3d9547b7215be034f04a1040b9b51a71cef9cf Mon Sep 17 00:00:00 2001
From: Sanjeev kumar <sanjeev.soft.dev@gmail.com>
Date: Mon, 1 Jan 2024 22:43:47 +0530
Subject: [PATCH] Improve code quality

---
 tests/client.test.js |  5 +++--
 tests/utils.js       | 11 +++++++++++
 2 files changed, 14 insertions(+), 2 deletions(-)

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..16869df 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'
+  };
+}