89 lines
2.5 KiB
TypeScript
89 lines
2.5 KiB
TypeScript
import { describe, it, expect, beforeEach, vi } from 'vitest'
|
|
import JanModelExtension from './index'
|
|
import ky from 'ky'
|
|
import { ModelManager } from '@janhq/core'
|
|
|
|
const API_URL = 'http://localhost:3000'
|
|
|
|
vi.stubGlobal('API_URL', API_URL)
|
|
|
|
describe('JanModelExtension', () => {
|
|
let extension: JanModelExtension
|
|
|
|
beforeEach(() => {
|
|
extension = new JanModelExtension()
|
|
vi.spyOn(ModelManager, 'instance').mockReturnValue({
|
|
get: (modelId: string) => ({
|
|
id: modelId,
|
|
engine: 'nitro_tensorrt_llm',
|
|
settings: { vision_model: true },
|
|
sources: [{ filename: 'test.bin' }],
|
|
}),
|
|
} as any)
|
|
vi.spyOn(JanModelExtension.prototype, 'cancelModelPull').mockImplementation(
|
|
async (model: string) => {
|
|
const kyDeleteSpy = vi.spyOn(ky, 'delete').mockResolvedValue({
|
|
json: () => Promise.resolve({}),
|
|
} as any)
|
|
|
|
await ky.delete(`${API_URL}/v1/models/pull`, {
|
|
json: { taskId: model },
|
|
})
|
|
|
|
expect(kyDeleteSpy).toHaveBeenCalledWith(`${API_URL}/v1/models/pull`, {
|
|
json: { taskId: model },
|
|
})
|
|
|
|
kyDeleteSpy.mockRestore() // Restore the original implementation
|
|
}
|
|
)
|
|
})
|
|
|
|
it('should initialize with an empty queue', () => {
|
|
expect(extension.queue.size).toBe(0)
|
|
})
|
|
|
|
describe('pullModel', () => {
|
|
it('should call the pull model endpoint with correct parameters', async () => {
|
|
const model = 'test-model'
|
|
const id = 'test-id'
|
|
const name = 'test-name'
|
|
|
|
const kyPostSpy = vi.spyOn(ky, 'post').mockReturnValue({
|
|
json: () => Promise.resolve({}),
|
|
} as any)
|
|
|
|
await extension.pullModel(model, id, name)
|
|
|
|
expect(kyPostSpy).toHaveBeenCalledWith(`${API_URL}/v1/models/pull`, {
|
|
json: { model, id, name },
|
|
})
|
|
|
|
kyPostSpy.mockRestore() // Restore the original implementation
|
|
})
|
|
})
|
|
|
|
describe('cancelModelPull', () => {
|
|
it('should call the cancel model pull endpoint with the correct model', async () => {
|
|
const model = 'test-model'
|
|
|
|
await extension.cancelModelPull(model)
|
|
})
|
|
})
|
|
|
|
describe('deleteModel', () => {
|
|
it('should call the delete model endpoint with the correct model', async () => {
|
|
const model = 'test-model'
|
|
const kyDeleteSpy = vi
|
|
.spyOn(ky, 'delete')
|
|
.mockResolvedValue({ json: () => Promise.resolve({}) } as any)
|
|
|
|
await extension.deleteModel(model)
|
|
|
|
expect(kyDeleteSpy).toHaveBeenCalledWith(`${API_URL}/v1/models/${model}`)
|
|
|
|
kyDeleteSpy.mockRestore() // Restore the original implementation
|
|
})
|
|
})
|
|
})
|