mirror of
https://github.com/Comfy-Org/ComfyUI_frontend.git
synced 2026-04-26 01:09:46 +00:00
[feat] Implement AssetService behind settings flag (#5404)
* [feat] add Comfy.Assets.UseAssetAPI to CORE_SETTINGS * [feat] create AssetService 1. Add service for accessing new Asset API 2. Add fallback model paths logic so empty model directories appear for the user. 3. Copious tests for them all. Co-Authored-By: Claude <noreply@anthropic.com> * [feat] switch between assets and file paths for model data * [feat] ignore assets with "missing" tag * [fix] formatting and style * [fix] call assets API with the correct filters * [feat] elminate unused modelPath code * [fix] remove stray comment * [fix] model manager api was not parsed correctly --------- Co-authored-by: Claude <noreply@anthropic.com>
This commit is contained in:
@@ -2,7 +2,9 @@ import { createPinia, setActivePinia } from 'pinia'
|
||||
import { beforeEach, describe, expect, it, vi } from 'vitest'
|
||||
|
||||
import { api } from '@/scripts/api'
|
||||
import { assetService } from '@/services/assetService'
|
||||
import { useModelStore } from '@/stores/modelStore'
|
||||
import { useSettingStore } from '@/stores/settingStore'
|
||||
|
||||
// Mock the api
|
||||
vi.mock('@/scripts/api', () => ({
|
||||
@@ -13,7 +15,32 @@ vi.mock('@/scripts/api', () => ({
|
||||
}
|
||||
}))
|
||||
|
||||
function enableMocks() {
|
||||
// Mock the assetService
|
||||
vi.mock('@/services/assetService', () => ({
|
||||
assetService: {
|
||||
getAssetModelFolders: vi.fn(),
|
||||
getAssetModels: vi.fn()
|
||||
}
|
||||
}))
|
||||
|
||||
// Mock the settingStore
|
||||
vi.mock('@/stores/settingStore', () => ({
|
||||
useSettingStore: vi.fn()
|
||||
}))
|
||||
|
||||
function enableMocks(useAssetAPI = false) {
|
||||
// Mock settingStore to return the useAssetAPI setting
|
||||
const mockSettingStore = {
|
||||
get: vi.fn().mockImplementation((key: string) => {
|
||||
if (key === 'Comfy.Assets.UseAssetAPI') {
|
||||
return useAssetAPI
|
||||
}
|
||||
return false
|
||||
})
|
||||
}
|
||||
vi.mocked(useSettingStore).mockReturnValue(mockSettingStore as any)
|
||||
|
||||
// Mock experimental API - returns objects with name and folders properties
|
||||
vi.mocked(api.getModels).mockResolvedValue([
|
||||
{ name: 'sdxl.safetensors', pathIndex: 0 },
|
||||
{ name: 'sdv15.safetensors', pathIndex: 0 },
|
||||
@@ -23,6 +50,18 @@ function enableMocks() {
|
||||
{ name: 'checkpoints', folders: ['/path/to/checkpoints'] },
|
||||
{ name: 'vae', folders: ['/path/to/vae'] }
|
||||
])
|
||||
|
||||
// Mock asset API - also returns objects with name and folders properties
|
||||
vi.mocked(assetService.getAssetModelFolders).mockResolvedValue([
|
||||
{ name: 'checkpoints', folders: ['/path/to/checkpoints'] },
|
||||
{ name: 'vae', folders: ['/path/to/vae'] }
|
||||
])
|
||||
vi.mocked(assetService.getAssetModels).mockResolvedValue([
|
||||
{ name: 'sdxl.safetensors', pathIndex: 0 },
|
||||
{ name: 'sdv15.safetensors', pathIndex: 0 },
|
||||
{ name: 'noinfo.safetensors', pathIndex: 0 }
|
||||
])
|
||||
|
||||
vi.mocked(api.viewMetadata).mockImplementation((_, model) => {
|
||||
if (model === 'noinfo.safetensors') {
|
||||
return Promise.resolve({})
|
||||
@@ -46,26 +85,25 @@ describe('useModelStore', () => {
|
||||
|
||||
beforeEach(async () => {
|
||||
setActivePinia(createPinia())
|
||||
store = useModelStore()
|
||||
vi.resetAllMocks()
|
||||
})
|
||||
|
||||
it('should load models', async () => {
|
||||
enableMocks()
|
||||
store = useModelStore()
|
||||
await store.loadModelFolders()
|
||||
const folderStore = await store.getLoadedModelFolder('checkpoints')
|
||||
expect(folderStore).not.toBeNull()
|
||||
if (!folderStore) return
|
||||
expect(Object.keys(folderStore.models).length).toBe(3)
|
||||
expect(folderStore).toBeDefined()
|
||||
expect(Object.keys(folderStore!.models)).toHaveLength(3)
|
||||
})
|
||||
|
||||
it('should load model metadata', async () => {
|
||||
enableMocks()
|
||||
store = useModelStore()
|
||||
await store.loadModelFolders()
|
||||
const folderStore = await store.getLoadedModelFolder('checkpoints')
|
||||
expect(folderStore).not.toBeNull()
|
||||
if (!folderStore) return
|
||||
const model = folderStore.models['0/sdxl.safetensors']
|
||||
expect(folderStore).toBeDefined()
|
||||
const model = folderStore!.models['0/sdxl.safetensors']
|
||||
await model.load()
|
||||
expect(model.title).toBe('Title of sdxl.safetensors')
|
||||
expect(model.architecture_id).toBe('stable-diffusion-xl-base-v1')
|
||||
@@ -79,11 +117,11 @@ describe('useModelStore', () => {
|
||||
|
||||
it('should handle no metadata', async () => {
|
||||
enableMocks()
|
||||
store = useModelStore()
|
||||
await store.loadModelFolders()
|
||||
const folderStore = await store.getLoadedModelFolder('checkpoints')
|
||||
expect(folderStore).not.toBeNull()
|
||||
if (!folderStore) return
|
||||
const model = folderStore.models['0/noinfo.safetensors']
|
||||
expect(folderStore).toBeDefined()
|
||||
const model = folderStore!.models['0/noinfo.safetensors']
|
||||
await model.load()
|
||||
expect(model.file_name).toBe('noinfo.safetensors')
|
||||
expect(model.title).toBe('noinfo')
|
||||
@@ -95,6 +133,7 @@ describe('useModelStore', () => {
|
||||
|
||||
it('should cache model information', async () => {
|
||||
enableMocks()
|
||||
store = useModelStore()
|
||||
await store.loadModelFolders()
|
||||
expect(api.getModels).toHaveBeenCalledTimes(0)
|
||||
await store.getLoadedModelFolder('checkpoints')
|
||||
@@ -102,4 +141,36 @@ describe('useModelStore', () => {
|
||||
await store.getLoadedModelFolder('checkpoints')
|
||||
expect(api.getModels).toHaveBeenCalledTimes(1)
|
||||
})
|
||||
|
||||
describe('API switching functionality', () => {
|
||||
it('should use experimental API for complete workflow when UseAssetAPI setting is false', async () => {
|
||||
enableMocks(false) // useAssetAPI = false
|
||||
store = useModelStore()
|
||||
await store.loadModelFolders()
|
||||
const folderStore = await store.getLoadedModelFolder('checkpoints')
|
||||
|
||||
// Both APIs return objects with .name property, modelStore extracts folder.name in both cases
|
||||
expect(api.getModelFolders).toHaveBeenCalledTimes(1)
|
||||
expect(api.getModels).toHaveBeenCalledWith('checkpoints')
|
||||
expect(assetService.getAssetModelFolders).toHaveBeenCalledTimes(0)
|
||||
expect(assetService.getAssetModels).toHaveBeenCalledTimes(0)
|
||||
expect(folderStore).toBeDefined()
|
||||
expect(Object.keys(folderStore!.models)).toHaveLength(3)
|
||||
})
|
||||
|
||||
it('should use asset API for complete workflow when UseAssetAPI setting is true', async () => {
|
||||
enableMocks(true) // useAssetAPI = true
|
||||
store = useModelStore()
|
||||
await store.loadModelFolders()
|
||||
const folderStore = await store.getLoadedModelFolder('checkpoints')
|
||||
|
||||
// Both APIs return objects with .name property, modelStore extracts folder.name in both cases
|
||||
expect(assetService.getAssetModelFolders).toHaveBeenCalledTimes(1)
|
||||
expect(assetService.getAssetModels).toHaveBeenCalledWith('checkpoints')
|
||||
expect(api.getModelFolders).toHaveBeenCalledTimes(0)
|
||||
expect(api.getModels).toHaveBeenCalledTimes(0)
|
||||
expect(folderStore).toBeDefined()
|
||||
expect(Object.keys(folderStore!.models)).toHaveLength(3)
|
||||
})
|
||||
})
|
||||
})
|
||||
|
||||
Reference in New Issue
Block a user