From 83edc1fbc729aa1ec4e5dd24a5385d8b49e3aa7c Mon Sep 17 00:00:00 2001
From: Louis <louis@jan.ai>
Date: Tue, 29 Oct 2024 14:30:21 +0700
Subject: [PATCH] chore: linter and test

---
 .../extensions/engines/LocalOAIEngine.test.ts | 43 ++++++++++++++++---
 web/containers/Providers/EventListener.tsx    |  3 +-
 2 files changed, 39 insertions(+), 7 deletions(-)

diff --git a/core/src/browser/extensions/engines/LocalOAIEngine.test.ts b/core/src/browser/extensions/engines/LocalOAIEngine.test.ts
index 8a7722f3a..e4296468f 100644
--- a/core/src/browser/extensions/engines/LocalOAIEngine.test.ts
+++ b/core/src/browser/extensions/engines/LocalOAIEngine.test.ts
@@ -43,22 +43,55 @@ describe('LocalOAIEngine', () => {
   })
 
   it('should load model correctly', async () => {
-    const model: Model = { engine: 'testProvider', file_path: 'path/to/model' } as any
+    const model: ModelFile = { engine: 'testProvider', file_path: 'path/to/model' } as any
+    const modelFolder = 'path/to'
+    const systemInfo = { os: 'testOS' }
+    const res = { error: null }
 
-    expect(engine.loadModel(model)).toBeTruthy()
+    ;(dirName as jest.Mock).mockResolvedValue(modelFolder)
+    ;(systemInformation as jest.Mock).mockResolvedValue(systemInfo)
+    ;(executeOnMain as jest.Mock).mockResolvedValue(res)
+
+    await engine.loadModel(model)
+
+    expect(dirName).toHaveBeenCalledWith(model.file_path)
+    expect(systemInformation).toHaveBeenCalled()
+    expect(executeOnMain).toHaveBeenCalledWith(
+      engine.nodeModule,
+      engine.loadModelFunctionName,
+      { modelFolder, model },
+      systemInfo
+    )
+    expect(events.emit).toHaveBeenCalledWith(ModelEvent.OnModelReady, model)
+  })
+
+  it('should handle load model error', async () => {
+    const model: any = { engine: 'testProvider', file_path: 'path/to/model' } as any
+    const modelFolder = 'path/to'
+    const systemInfo = { os: 'testOS' }
+    const res = { error: 'load error' }
+
+    ;(dirName as jest.Mock).mockResolvedValue(modelFolder)
+    ;(systemInformation as jest.Mock).mockResolvedValue(systemInfo)
+    ;(executeOnMain as jest.Mock).mockResolvedValue(res)
+
+    await expect(engine.loadModel(model)).rejects.toEqual('load error')
+
+    expect(events.emit).toHaveBeenCalledWith(ModelEvent.OnModelFail, { error: res.error })
   })
 
   it('should unload model correctly', async () => {
     const model: Model = { engine: 'testProvider' } as any
 
-    expect(engine.unloadModel(model)).toBeTruthy()
+    await engine.unloadModel(model)
+
+    expect(executeOnMain).toHaveBeenCalledWith(engine.nodeModule, engine.unloadModelFunctionName)
+    expect(events.emit).toHaveBeenCalledWith(ModelEvent.OnModelStopped, {})
   })
 
   it('should not unload model if engine does not match', async () => {
     const model: Model = { engine: 'otherProvider' } as any
-
     await engine.unloadModel(model)
-
     expect(executeOnMain).not.toHaveBeenCalled()
     expect(events.emit).not.toHaveBeenCalledWith(ModelEvent.OnModelStopped, {})
   })
diff --git a/web/containers/Providers/EventListener.tsx b/web/containers/Providers/EventListener.tsx
index af91b6027..9535bbfa6 100644
--- a/web/containers/Providers/EventListener.tsx
+++ b/web/containers/Providers/EventListener.tsx
@@ -88,8 +88,7 @@ const EventListenerWrapper = ({ children }: PropsWithChildren) => {
       if (state.downloadType !== 'extension') {
         state.downloadState = 'end'
         setDownloadState(state)
-        if (state.percent !== 0)
-          removeDownloadingModel(state.modelId)
+        if (state.percent !== 0) removeDownloadingModel(state.modelId)
       }
       events.emit(ModelEvent.OnModelsUpdate, {})
     },