fix: use spy for openaiclient class

opengovsg · Mar 6, 2024 · 4d6feb1 · 4d6feb1
1 parent 0514e77
commit 4d6feb1
Showing 1 changed file with 7 additions and 24 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
@@ -15,29 +15,7 @@ const mockReturnValue = {
 }
 
 // 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)
+jest.mock('@azure/openai')
 
 beforeEach(() => {
   jest.clearAllMocks()
@@ -49,14 +27,19 @@ describe('admin-form.assistance.service', () => {
       // Arrange
       const type = 'prompt'
       const content = 'Sample content'
+      const mockReturnMessage = 'some message'
+
+      jest
+        .spyOn(OpenAIClient.prototype, 'getChatCompletions')
+        .mockResolvedValue({ choices: [{ message: mockReturnMessage }] } as any)
 
       // Act
       const actualResult = await generateQuestions({ type, content })
 
       // Assert
       expect(actualResult.isOk()).toEqual(true)
 
-      expect(actualResult._unsafeUnwrap()).toMatchObject(mockReturnValue)
+      expect(actualResult._unsafeUnwrap()).toMatch(mockReturnMessage)
     })
 
     it('should return list of questions based on the PDF type and content', async () => {