Files
ComfyUI_frontend/tests-ui/tests/store/modelStore.test.ts
Alex "mcmonkey" Goodwin 060e61f0db initial model store (#674)
* initial model store

* refactor the 'modelstoreserviceimpl' to pinia

* pepper in some reactive

the inner ModelStore (per-folder) can't be pinia because its made of temporary instances, but it can be reactive

* use refs in metadata class

* remove 'reactive'

* remove ref too

* add simple unit tests for modelStore

* make things worse via autoformatting

* move mock impls to a function
2024-09-04 19:59:40 -04:00

86 lines
2.8 KiB
TypeScript

import { setActivePinia, createPinia } from 'pinia'
import { useModelStore } from '@/stores/modelStore'
import { api } from '@/scripts/api'
// Mock the api
jest.mock('@/scripts/api', () => ({
api: {
getModels: jest.fn(),
viewMetadata: jest.fn()
}
}))
function enableMocks() {
;(api.getModels as jest.Mock).mockResolvedValue([
'sdxl.safetensors',
'sdv15.safetensors',
'noinfo.safetensors'
])
;(api.viewMetadata as jest.Mock).mockImplementation((_, model) => {
if (model === 'noinfo.safetensors') {
return Promise.resolve({})
}
return Promise.resolve({
'modelspec.title': `Title of ${model}`,
display_name: 'Should not show',
'modelspec.architecture': 'stable-diffusion-xl-base-v1',
'modelspec.author': `Author of ${model}`,
'modelspec.description': `Description of ${model}`,
'modelspec.resolution': '1024x1024',
trigger_phrase: `Trigger phrase of ${model}`,
usage_hint: `Usage hint of ${model}`,
tags: `tags,for,${model}`
})
})
}
describe('useModelStore', () => {
let store: ReturnType<typeof useModelStore>
beforeEach(() => {
setActivePinia(createPinia())
store = useModelStore()
})
it('should load models', async () => {
enableMocks()
const folderStore = await store.getModelsInFolderCached('checkpoints')
expect(folderStore).not.toBeNull()
expect(Object.keys(folderStore.models)).toHaveLength(3)
})
it('should load model metadata', async () => {
enableMocks()
const folderStore = await store.getModelsInFolderCached('checkpoints')
const model = folderStore.models['sdxl.safetensors']
await model.load()
expect(model.title).toBe('Title of sdxl.safetensors')
expect(model.architecture_id).toBe('stable-diffusion-xl-base-v1')
expect(model.author).toBe('Author of sdxl.safetensors')
expect(model.description).toBe('Description of sdxl.safetensors')
expect(model.resolution).toBe('1024x1024')
expect(model.trigger_phrase).toBe('Trigger phrase of sdxl.safetensors')
expect(model.usage_hint).toBe('Usage hint of sdxl.safetensors')
expect(model.tags).toHaveLength(3)
})
it('should handle no metadata', async () => {
enableMocks()
const folderStore = await store.getModelsInFolderCached('checkpoints')
const model = folderStore.models['noinfo.safetensors']
await model.load()
expect(model.title).toBe('noinfo.safetensors')
expect(model.architecture_id).toBe('')
expect(model.author).toBe('')
expect(model.description).toBe('')
expect(model.resolution).toBe('')
})
it('should cache model information', async () => {
enableMocks()
const folderStore1 = await store.getModelsInFolderCached('checkpoints')
const folderStore2 = await store.getModelsInFolderCached('checkpoints')
expect(api.getModels).toHaveBeenCalledTimes(1)
})
})