Fix tests

SeoulSKY · Apr 26, 2024 · 2aaee0f · 2aaee0f
1 parent 202f3db
commit 2aaee0f
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 15 deletions.
diff --git a/__tests__/index.ts b/__tests__/index.ts
@@ -1,5 +1,22 @@
 import {InvalidArgumentError} from "../src/utils/errors";
 import {Storage} from "../src/utils/storage";
+import {spyOn} from "jest-mock";
+import {ChatSession, FunctionCall} from "@google/generative-ai";
+
+export function mockChatSession() {
+  return spyOn(ChatSession.prototype, "sendMessage")
+    .mockResolvedValue({
+      response: {
+        text: () => "response",
+        functionCalls: function (): FunctionCall[] | undefined {
+          throw new Error("Function not implemented.");
+        },
+        functionCall: function (): FunctionCall | undefined {
+          throw new Error("Function not implemented.");
+        }
+      }
+    });
+}
 
 export class MockStorage<K, V> implements Storage<K, V> {
 

diff --git a/__tests__/utils/chat.test.ts b/__tests__/utils/chat.test.ts
@@ -1,13 +1,12 @@
 import {beforeAll, beforeEach, describe, expect, it} from "@jest/globals";
 import {spyOn} from "jest-mock";
 
-import {MockStorage} from "../index";
+import {mockChatSession, MockStorage} from "../index";
 
 import Chat, {Message} from "../../src/utils/chat";
 import {InvalidArgumentError, InvalidStateError} from "../../src/utils/errors";
 import {UserProfile, BotProfile, Participant, Gender} from "../../src/utils/profile";
 import {ImageData} from "../../src/utils/image";
-import { FunctionCall } from "@google/generative-ai";
 
 const mockProfileData = {
   name: "test",
@@ -42,23 +41,13 @@ describe("Chat", () => {
     await UserProfile.getInstance(MockStorage).create(mockProfileData);
     await BotProfile.getInstance(MockStorage).create(mockProfileData);
 
+    mockChatSession();
+
     chat = await Chat.getInstance(MockStorage);
 
     storage = new MockStorage();
     chat["storage"] = storage;
 
-    spyOn(chat["session"], "sendMessage")
-      .mockResolvedValue({
-        response: {
-          text: () => "response",
-          functionCalls: function (): FunctionCall[] | undefined {
-            throw new Error("Function not implemented.");
-          },
-          functionCall: function (): FunctionCall | undefined {
-            throw new Error("Function not implemented.");
-          }
-        }});
-
     // @ts-expect-error for testing purposes
     spyOn(chat, "getImageDescriptions").mockResolvedValue(Promise.resolve("image descriptions"));
   });

diff --git a/__tests__/utils/quiz.test.ts b/__tests__/utils/quiz.test.ts
@@ -1,7 +1,7 @@
 import {beforeAll, beforeEach, describe, expect, it} from "@jest/globals";
 import {spyOn} from "jest-mock";
 
-import {MockStorage} from "../index";
+import {mockChatSession, MockStorage} from "../index";
 
 import Quiz, {Difficulty, MultipleChoiceQuestion} from "../../src/utils/quiz";
 import Chat, {Message} from "../../src/utils/chat";
@@ -129,6 +129,8 @@ describe("Quiz", () => {
     await BotProfile.getInstance(MockStorage).create(mockProfileData);
     await UserProfile.getInstance(MockStorage).create(mockProfileData);
 
+    mockChatSession();
+
     chat = await Chat.getInstance(MockStorage);
   });