From 72e9828b763d1071a2b48ea81cd432dc2be65600 Mon Sep 17 00:00:00 2001
From: Danny Avila <110412045+danny-avila@users.noreply.github.com>
Date: Tue, 13 Jun 2023 00:04:01 -0400
Subject: [PATCH] tests(api): refactor to mock database and network operations
 (#494)

---
 api/app/langchain/ChatAgent.test.js   | 83 ++++++++++++++++++++++-----
 api/app/langchain/tools/index.test.js | 47 ++++++++++-----
 2 files changed, 103 insertions(+), 27 deletions(-)

diff --git a/api/app/langchain/ChatAgent.test.js b/api/app/langchain/ChatAgent.test.js
index 16b37575b..ce18b98ff 100644
--- a/api/app/langchain/ChatAgent.test.js
+++ b/api/app/langchain/ChatAgent.test.js
@@ -1,7 +1,16 @@
-const mongoose = require('mongoose');
+const { HumanChatMessage, AIChatMessage } = require('langchain/schema');
 const ChatAgent = require('./ChatAgent');
-const connectDb = require('../../lib/db/connectDb');
-const Conversation = require('../../models/Conversation');
+const crypto = require('crypto');
+
+jest.mock('../../lib/db/connectDb');
+jest.mock('../../models/Conversation', () => {
+  return function () {
+    return {
+      save: jest.fn(),
+      deleteConvos: jest.fn()
+    };
+  };
+});
 
 describe('ChatAgent', () => {
   let TestAgent;
@@ -13,26 +22,72 @@ describe('ChatAgent', () => {
       max_tokens: 2
     },
     agentOptions: {
-      model: 'gpt-3.5-turbo',
+      model: 'gpt-3.5-turbo'
     }
   };
   let parentMessageId;
   let conversationId;
+  const fakeMessages = [];
   const userMessage = 'Hello, ChatGPT!';
-  const apiKey = process.env.OPENAI_API_KEY;
-
-  beforeAll(async () => {
-    await connectDb();
-  });
+  const apiKey = 'fake-api-key';
 
   beforeEach(() => {
     TestAgent = new ChatAgent(apiKey, options);
-  });
+    TestAgent.loadHistory = jest
+      .fn()
+      .mockImplementation((conversationId, parentMessageId = null) => {
+        if (!conversationId) {
+          TestAgent.currentMessages = [];
+          return Promise.resolve([]);
+        }
 
-  afterAll(async () => {
-    // Delete the messages and conversation created by the test
-    await Conversation.deleteConvos(null, { conversationId });
-    await mongoose.connection.close();
+        const orderedMessages = TestAgent.constructor.getMessagesForConversation(
+          fakeMessages,
+          parentMessageId
+        );
+        const chatMessages = orderedMessages.map((msg) =>
+          msg?.isCreatedByUser || msg?.role.toLowerCase() === 'user'
+            ? new HumanChatMessage(msg.text)
+            : new AIChatMessage(msg.text)
+        );
+
+        TestAgent.currentMessages = orderedMessages;
+        return Promise.resolve(chatMessages);
+      });
+    TestAgent.sendMessage = jest.fn().mockImplementation(async (message, opts = {}) => {
+      if (opts && typeof opts === 'object') {
+        TestAgent.setOptions(opts);
+      }
+      const conversationId = opts.conversationId || crypto.randomUUID();
+      const parentMessageId = opts.parentMessageId || '00000000-0000-0000-0000-000000000000';
+      const userMessageId = opts.overrideParentMessageId || crypto.randomUUID();
+      this.pastMessages = await TestAgent.loadHistory(
+        conversationId,
+        TestAgent.options?.parentMessageId
+      );
+
+      const userMessage = {
+        text: message,
+        sender: 'ChatGPT',
+        isCreatedByUser: true,
+        messageId: userMessageId,
+        parentMessageId,
+        conversationId
+      };
+
+      const response = {
+        sender: 'ChatGPT',
+        text: 'Hello, User!',
+        isCreatedByUser: false,
+        messageId: crypto.randomUUID(),
+        parentMessageId: userMessage.messageId,
+        conversationId
+      };
+
+      fakeMessages.push(userMessage);
+      fakeMessages.push(response);
+      return response;
+    });
   });
 
   test('initializes ChatAgent without crashing', () => {
diff --git a/api/app/langchain/tools/index.test.js b/api/app/langchain/tools/index.test.js
index e2917cad0..9cd9ccd15 100644
--- a/api/app/langchain/tools/index.test.js
+++ b/api/app/langchain/tools/index.test.js
@@ -1,8 +1,25 @@
-/* eslint-disable jest/no-conditional-expect */
-require('dotenv').config({ path: '../../../.env' });
-const mongoose = require('mongoose');
+const mockUser = {
+  _id: 'fakeId',
+  save: jest.fn(),
+  findByIdAndDelete: jest.fn(),
+};
+
+var mockPluginService = {
+  updateUserPluginAuth: jest.fn(),
+  deleteUserPluginAuth: jest.fn(),
+  getUserPluginAuthValue: jest.fn()
+};
+
+
+jest.mock('../../../models/User', () => {
+  return function() {
+    return mockUser;
+  };
+});
+
+jest.mock('../../../server/services/PluginService', () => mockPluginService);
+
 const User = require('../../../models/User');
-const connectDb = require('../../../lib/db/connectDb');
 const { validateTools, loadTools, availableTools } = require('./index');
 const PluginService = require('../../../server/services/PluginService');
 const { BaseChatModel } = require('langchain/chat_models/openai');
@@ -21,7 +38,16 @@ describe('Tool Handlers', () => {
   const authConfigs = mainPlugin.authConfig;
 
   beforeAll(async () => {
-    await connectDb();
+    mockUser.save.mockResolvedValue(undefined);
+  
+    const userAuthValues = {};
+    mockPluginService.getUserPluginAuthValue.mockImplementation((userId, authField) => {
+      return userAuthValues[`${userId}-${authField}`];
+    });
+    mockPluginService.updateUserPluginAuth.mockImplementation((userId, authField, _pluginKey, credential) => {
+      userAuthValues[`${userId}-${authField}`] = credential;
+    });
+  
     fakeUser = new User({
       name: 'Fake User',
       username: 'fakeuser',
@@ -39,19 +65,13 @@ describe('Tool Handlers', () => {
     for (const authConfig of authConfigs) {
       await PluginService.updateUserPluginAuth(fakeUser._id, authConfig.authField, pluginKey, mockCredential);
     }
-  });
-
-  // afterEach(async () => {
-  //   // Clean up any test-specific data.
-  // });
+  });  
 
   afterAll(async () => {
-    // Delete the fake user & plugin auth
-    await User.findByIdAndDelete(fakeUser._id);
+    await mockUser.findByIdAndDelete(fakeUser._id);
     for (const authConfig of authConfigs) {
       await PluginService.deleteUserPluginAuth(fakeUser._id, authConfig.authField);
     }
-    await mongoose.connection.close();
   });
 
   describe('validateTools', () => {
@@ -128,6 +148,7 @@ describe('Tool Handlers', () => {
       try {
         await loadTool2();
       } catch (error) {
+        // eslint-disable-next-line jest/no-conditional-expect
         expect(error).toBeDefined();
       }
     });