Skip to content

Commit

Permalink
test fixes
Browse files Browse the repository at this point in the history
  • Loading branch information
nbonamy committed Jun 29, 2024
1 parent b12d349 commit 4a0f084
Show file tree
Hide file tree
Showing 3 changed files with 25 additions and 23 deletions.
2 changes: 1 addition & 1 deletion src/services/ollama.ts
Original file line number Diff line number Diff line change
Expand Up @@ -95,7 +95,7 @@ export default class extends LlmEngine {
}

async stop() {
await ollama.abort()
await this.client.abort()
}

// eslint-disable-next-line @typescript-eslint/no-unused-vars
Expand Down
2 changes: 2 additions & 0 deletions tests/unit/config.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,8 @@ vi.mock('fs', async (importOriginal) => {
test('Load default settings', () => {
const loaded = config.loadSettings('')
delete loaded.getActiveModel
loaded.engines.openai.baseURL = defaultSettings.engines.openai.baseURL
loaded.engines.ollama.baseURL = defaultSettings.engines.ollama.baseURL
expect(loaded).toStrictEqual(defaultSettings)
expect(loaded.general.language).toBe('')
expect(loaded.engines.openai.models.chat).toStrictEqual([])
Expand Down
44 changes: 22 additions & 22 deletions tests/unit/engine_ollama.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -10,27 +10,27 @@ import { loadOllamaModels } from '../../src/services/llm'
import { Model } from '../../src/types/config.d'

vi.mock('ollama/browser', async() => {
return { default : {
list: vi.fn(() => {
return { models: [
{ model: 'model2', name: 'model2' },
{ model: 'model1', name: 'model1' },
] }
}),
chat: vi.fn((opts) => {
if (opts.stream) {
return {
controller: {
abort: vi.fn()
}
const Ollama = vi.fn()
Ollama.prototype.list = vi.fn(() => {
return { models: [
{ model: 'model2', name: 'model2' },
{ model: 'model1', name: 'model1' },
] }
})
Ollama.prototype.chat = vi.fn((opts) => {
if (opts.stream) {
return {
controller: {
abort: vi.fn()
}
}
else {
return { message: { content: 'response' } }
}
}),
abort: vi.fn(),
}}
}
else {
return { message: { content: 'response' } }
}
})
Ollama.prototype.abort = vi.fn()
return { Ollama: Ollama }
})

beforeEach(() => {
Expand Down Expand Up @@ -61,7 +61,7 @@ test('Ollama completion', async () => {
new Message('system', 'instruction'),
new Message('user', 'prompt'),
], null)
expect(_ollama.default.chat).toHaveBeenCalled()
expect(_ollama.Ollama.prototype.chat).toHaveBeenCalled()
expect(response).toStrictEqual({
type: 'text',
content: 'response'
Expand All @@ -74,10 +74,10 @@ test('Ollama stream', async () => {
new Message('system', 'instruction'),
new Message('user', 'prompt'),
], null)
expect(_ollama.default.chat).toHaveBeenCalled()
expect(_ollama.Ollama.prototype.chat).toHaveBeenCalled()
expect(response.controller).toBeDefined()
await ollama.stop()
expect(_ollama.default.abort).toHaveBeenCalled()
expect(_ollama.Ollama.prototype.abort).toHaveBeenCalled()
})

test('Ollama image', async () => {
Expand Down

0 comments on commit 4a0f084

Please sign in to comment.