From aa474663bbbfff876aa05a3ca453840b6600d728 Mon Sep 17 00:00:00 2001
From: Kathleen Koh <kathleen@open.gov.sg>
Date: Mon, 4 Mar 2024 23:29:56 +0800
Subject: [PATCH] chore: add backend tests wip

---
 .../admin-form.assistance.service.spec.ts     | 48 +++++++++++--------
 .../admin-form.assistance.service.ts          |  8 +---
 2 files changed, 30 insertions(+), 26 deletions(-)

diff --git a/src/app/modules/form/admin-form/__tests__/admin-form.assistance.service.spec.ts b/src/app/modules/form/admin-form/__tests__/admin-form.assistance.service.spec.ts
index 2f55bf57df..14b50c9fb1 100644
--- a/src/app/modules/form/admin-form/__tests__/admin-form.assistance.service.spec.ts
+++ b/src/app/modules/form/admin-form/__tests__/admin-form.assistance.service.spec.ts
@@ -1,4 +1,4 @@
-import OpenAI from 'openai'
+import { OpenAIClient } from '@azure/openai'
 
 import {
   generateFormFields,
@@ -9,28 +9,38 @@ import {
   AssistanceModelTypeError,
 } from '../admin-form.errors'
 
-// Mock openai
-jest.mock('openai', () => jest.fn())
-const MockedOpenAIClient = jest.mocked(OpenAI)
-
 const mockReturnValue = {
   role: 'user',
   content: 'dummy content',
 }
 
+// Mock azure openai
+jest.mock('@azure/openai', () => {
+  return {
+    AzureKeyCredential: jest.fn().mockImplementation((apiKey) => {
+      return apiKey
+    }),
+    OpenAIClient: jest.fn().mockImplementation(() => {
+      return {
+        getChatCompletions: jest.fn().mockResolvedValue({
+          choices: [
+            {
+              message: {
+                role: 'user',
+                content: 'dummy content',
+              },
+            },
+          ],
+        }),
+      }
+    }),
+  }
+})
+
+const MockedOpenAIClient = jest.mocked(OpenAIClient)
+
 beforeEach(() => {
   jest.clearAllMocks()
-  MockedOpenAIClient.prototype.chat = {
-    completions: {
-      create: jest.fn().mockResolvedValue({
-        choices: [
-          {
-            message: mockReturnValue,
-          },
-        ],
-      }),
-    },
-  } as any
 })
 
 describe('admin-form.assistance.service', () => {
@@ -67,8 +77,8 @@ describe('admin-form.assistance.service', () => {
       const type = 'prompt'
       const content = 'Sample content'
 
-      // Mock OpenAI API throwing an error
-      MockedOpenAIClient.prototype.chat.completions.create = jest
+      // Mock Azure OpenAI API throwing an error
+      MockedOpenAIClient.prototype.getChatCompletions = jest
         .fn()
         .mockRejectedValue(new Error('Some random error message'))
 
@@ -119,7 +129,7 @@ describe('admin-form.assistance.service', () => {
       const questions = 'sample questions'
 
       // Mock OpenAI API throwing an error
-      MockedOpenAIClient.prototype.chat.completions.create = jest
+      MockedOpenAIClient.prototype.getChatCompletions = jest
         .fn()
         .mockRejectedValue(new Error('Some random error message'))
 
diff --git a/src/app/modules/form/admin-form/admin-form.assistance.service.ts b/src/app/modules/form/admin-form/admin-form.assistance.service.ts
index 41def3626a..02ba5fa4bb 100644
--- a/src/app/modules/form/admin-form/admin-form.assistance.service.ts
+++ b/src/app/modules/form/admin-form/admin-form.assistance.service.ts
@@ -32,13 +32,7 @@ const endpoint = azureOpenAIConfig.endpoint
 const azureApiKey = azureOpenAIConfig.apiKey
 const deploymentId = azureOpenAIConfig.deploymentId
 
-if (!endpoint || !azureApiKey) {
-  throw new Error(
-    'Azure OpenAI endpoint or API key is not defined in the environment variables',
-  )
-}
-
-const azureOpenAi = new OpenAIClient(
+export const azureOpenAi = new OpenAIClient(
   endpoint,
   new AzureKeyCredential(azureApiKey),
 )