fix: tests

core/src/browser/extensions/engines/LocalOAIEngine.test.ts

@@ -3,96 +3,134 @@
  */
 import { LocalOAIEngine } from './LocalOAIEngine'
 import { events } from '../../events'
-import { ModelEvent, Model } from '../../../types'
+import { Model, ModelEvent } from '../../../types'
-import { executeOnMain, systemInformation, dirName } from '../../core'
 
-jest.mock('../../core', () => ({
+jest.mock('../../events')
-  executeOnMain: jest.fn(),
-  systemInformation: jest.fn(),
-  dirName: jest.fn(),
-}))
-
-jest.mock('../../events', () => ({
-  events: {
-    on: jest.fn(),
-    emit: jest.fn(),
-  },
-}))
 
 class TestLocalOAIEngine extends LocalOAIEngine {
-  inferenceUrl = ''
+  inferenceUrl = 'http://test-local-inference-url'
-  nodeModule = 'testNodeModule'
+  provider = 'test-local-provider'
-  provider = 'testProvider'
+  nodeModule = 'test-node-module'
+
+  async headers() {
+    return { Authorization: 'Bearer test-token' }
+  }
+
+  async loadModel(model: Model & { file_path?: string }): Promise<void> {
+    this.loadedModel = model
+  }
+
+  async unloadModel(model?: Model) {
+    this.loadedModel = undefined
+  }
 }
 
 describe('LocalOAIEngine', () => {
   let engine: TestLocalOAIEngine
+  const mockModel: Model & { file_path?: string } = {
+    object: 'model',
+    version: '1.0.0',
+    format: 'gguf',
+    sources: [],
+    id: 'test-model',
+    name: 'Test Model',
+    description: 'A test model',
+    settings: {},
+    parameters: {},
+    metadata: {},
+    file_path: '/path/to/model.gguf'
+  }
 
   beforeEach(() => {
     engine = new TestLocalOAIEngine('', '')
-  })
-
-  afterEach(() => {
     jest.clearAllMocks()
   })
 
-  it('should subscribe to events on load', () => {
+  describe('onLoad', () => {
+    it('should call super.onLoad and subscribe to model events', () => {
+      const superOnLoadSpy = jest.spyOn(Object.getPrototypeOf(Object.getPrototypeOf(engine)), 'onLoad')
+      
       engine.onLoad()
-    expect(events.on).toHaveBeenCalledWith(ModelEvent.OnModelInit, expect.any(Function))
-    expect(events.on).toHaveBeenCalledWith(ModelEvent.OnModelStop, expect.any(Function))
-  })
 
-  it('should load model correctly', async () => {
+      expect(superOnLoadSpy).toHaveBeenCalled()
-    const model: any = { engine: 'testProvider', file_path: 'path/to/model' } as any
+      expect(events.on).toHaveBeenCalledWith(
-    const modelFolder = 'path/to'
+        ModelEvent.OnModelInit,
-    const systemInfo = { os: 'testOS' }
+        expect.any(Function)
-    const res = { error: null }
+      )
-
+      expect(events.on).toHaveBeenCalledWith(
-    ;(dirName as jest.Mock).mockResolvedValue(modelFolder)
+        ModelEvent.OnModelStop,
-    ;(systemInformation as jest.Mock).mockResolvedValue(systemInfo)
+        expect.any(Function)
-    ;(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 () => {
+    it('should load model when OnModelInit event is triggered', () => {
-    const model: any = { engine: 'testProvider', file_path: 'path/to/model' } as any
+      const loadModelSpy = jest.spyOn(engine, 'loadModel')
-    const modelFolder = 'path/to'
+      engine.onLoad()
-    const systemInfo = { os: 'testOS' }
-    const res = { error: 'load error' }
 
-    ;(dirName as jest.Mock).mockResolvedValue(modelFolder)
+      // Get the event handler for OnModelInit
-    ;(systemInformation as jest.Mock).mockResolvedValue(systemInfo)
+      const onModelInitCall = (events.on as jest.Mock).mock.calls.find(
-    ;(executeOnMain as jest.Mock).mockResolvedValue(res)
+        call => call[0] === ModelEvent.OnModelInit
+      )
+      const onModelInitHandler = onModelInitCall[1]
 
-    await expect(engine.loadModel(model)).rejects.toEqual('load error')
+      // Trigger the event handler
+      onModelInitHandler(mockModel)
 
-    expect(events.emit).toHaveBeenCalledWith(ModelEvent.OnModelFail, { error: res.error })
+      expect(loadModelSpy).toHaveBeenCalledWith(mockModel)
     })
 
-  it('should unload model correctly', async () => {
+    it('should unload model when OnModelStop event is triggered', () => {
-    const model: Model = { engine: 'testProvider' } as any
+      const unloadModelSpy = jest.spyOn(engine, 'unloadModel')
+      engine.onLoad()
 
-    await engine.unloadModel(model)
+      // Get the event handler for OnModelStop
+      const onModelStopCall = (events.on as jest.Mock).mock.calls.find(
+        call => call[0] === ModelEvent.OnModelStop
+      )
+      const onModelStopHandler = onModelStopCall[1]
 
-    expect(executeOnMain).toHaveBeenCalledWith(engine.nodeModule, engine.unloadModelFunctionName)
+      // Trigger the event handler
-    expect(events.emit).toHaveBeenCalledWith(ModelEvent.OnModelStopped, {})
+      onModelStopHandler(mockModel)
+
+      expect(unloadModelSpy).toHaveBeenCalledWith(mockModel)
+    })
   })
 
-  it('should not unload model if engine does not match', async () => {
+  describe('properties', () => {
-    const model: Model = { engine: 'otherProvider' } as any
+    it('should have correct default function names', () => {
-    await engine.unloadModel(model)
+      expect(engine.loadModelFunctionName).toBe('loadModel')
-    expect(executeOnMain).not.toHaveBeenCalled()
+      expect(engine.unloadModelFunctionName).toBe('unloadModel')
-    expect(events.emit).not.toHaveBeenCalledWith(ModelEvent.OnModelStopped, {})
+    })
+
+    it('should have abstract nodeModule property implemented', () => {
+      expect(engine.nodeModule).toBe('test-node-module')
+    })
+  })
+
+  describe('loadModel', () => {
+    it('should load the model and set loadedModel', async () => {
+      await engine.loadModel(mockModel)
+      expect(engine.loadedModel).toBe(mockModel)
+    })
+
+    it('should handle model with file_path', async () => {
+      const modelWithPath = { ...mockModel, file_path: '/custom/path/model.gguf' }
+      await engine.loadModel(modelWithPath)
+      expect(engine.loadedModel).toBe(modelWithPath)
+    })
+  })
+
+  describe('unloadModel', () => {
+    it('should unload the model and clear loadedModel', async () => {
+      engine.loadedModel = mockModel
+      await engine.unloadModel(mockModel)
+      expect(engine.loadedModel).toBeUndefined()
+    })
+
+    it('should handle unload without passing a model', async () => {
+      engine.loadedModel = mockModel
+      await engine.unloadModel()
+      expect(engine.loadedModel).toBeUndefined()
+    })
   })
 })