import { createPinia, setActivePinia } from 'pinia' import { beforeEach, describe, expect, it, vi } from 'vitest' import type { MissingModelCandidate } from '@/platform/missingModel/types' const mockGetNodeByExecutionId = vi.fn() const mockResolveNodeDisplayName = vi.fn() const mockValidateSourceUrl = vi.fn() const mockGetAssetMetadata = vi.fn() const mockGetAssetDisplayName = vi.fn((a: { name: string }) => a.name) const mockGetAssetFilename = vi.fn((a: { name: string }) => a.name) const mockGetAssets = vi.fn() const mockUpdateModelsForNodeType = vi.fn() const mockGetAllNodeProviders = vi.fn() const mockDownloadList = vi.fn( (): Array<{ taskId: string; status: string }> => [] ) vi.mock('@/i18n', () => ({ st: vi.fn((_key: string, fallback: string) => fallback) })) vi.mock('@/platform/distribution/types', () => ({ isCloud: false })) vi.mock('vue-i18n', () => ({ useI18n: () => ({ t: (key: string) => key }) })) vi.mock('@/scripts/app', () => ({ app: { rootGraph: null } })) vi.mock('@/utils/graphTraversalUtil', () => ({ getNodeByExecutionId: (...args: unknown[]) => mockGetNodeByExecutionId(...args) })) vi.mock('@/utils/nodeTitleUtil', () => ({ resolveNodeDisplayName: (...args: unknown[]) => mockResolveNodeDisplayName(...args) })) vi.mock('@/renderer/core/canvas/canvasStore', () => ({ useCanvasStore: () => ({}) })) vi.mock('@/stores/assetsStore', () => ({ useAssetsStore: () => ({ getAssets: mockGetAssets, updateModelsForNodeType: mockUpdateModelsForNodeType, invalidateModelsForCategory: vi.fn(), updateModelsForTag: vi.fn() }) })) vi.mock('@/stores/assetDownloadStore', () => ({ useAssetDownloadStore: () => ({ get downloadList() { return mockDownloadList() }, trackDownload: vi.fn() }) })) vi.mock('@/stores/modelToNodeStore', () => ({ useModelToNodeStore: () => ({ getAllNodeProviders: mockGetAllNodeProviders }) })) vi.mock('@/platform/assets/services/assetService', () => ({ assetService: { getAssetMetadata: (...args: unknown[]) => mockGetAssetMetadata(...args) } })) vi.mock('@/platform/assets/utils/assetMetadataUtils', () => ({ getAssetDisplayName: (a: { name: string }) => mockGetAssetDisplayName(a), getAssetFilename: (a: { name: string }) => mockGetAssetFilename(a) })) vi.mock('@/platform/assets/importSources/civitaiImportSource', () => ({ civitaiImportSource: { type: 'civitai', name: 'Civitai', hostnames: ['civitai.com'] } })) vi.mock('@/platform/assets/importSources/huggingfaceImportSource', () => ({ huggingfaceImportSource: { type: 'huggingface', name: 'Hugging Face', hostnames: ['huggingface.co'] } })) vi.mock('@/platform/assets/utils/importSourceUtil', () => ({ validateSourceUrl: (...args: unknown[]) => mockValidateSourceUrl(...args) })) import { app } from '@/scripts/app' import { useMissingModelStore } from '@/platform/missingModel/missingModelStore' import { getComboValue, getModelStateKey, getNodeDisplayLabel, useMissingModelInteractions } from './useMissingModelInteractions' function makeCandidate( overrides: Partial = {} ): MissingModelCandidate { return { name: 'model.safetensors', nodeId: '1', nodeType: 'CheckpointLoaderSimple', widgetName: 'ckpt_name', isAssetSupported: false, isMissing: true, ...overrides } } describe('useMissingModelInteractions', () => { beforeEach(() => { setActivePinia(createPinia()) vi.resetAllMocks() mockGetAssetDisplayName.mockImplementation((a: { name: string }) => a.name) mockGetAssetFilename.mockImplementation((a: { name: string }) => a.name) mockDownloadList.mockImplementation( (): Array<{ taskId: string; status: string }> => [] ) ;(app as { rootGraph: unknown }).rootGraph = null }) describe('getModelStateKey', () => { it('returns key with supported prefix when asset is supported', () => { expect(getModelStateKey('model.safetensors', 'checkpoints', true)).toBe( 'supported::checkpoints::model.safetensors' ) }) it('returns key with unsupported prefix when asset is not supported', () => { expect(getModelStateKey('model.safetensors', 'loras', false)).toBe( 'unsupported::loras::model.safetensors' ) }) it('handles null directory', () => { expect(getModelStateKey('model.safetensors', null, true)).toBe( 'supported::::model.safetensors' ) }) }) describe('getNodeDisplayLabel', () => { it('returns fallback when graph is null', () => { ;(app as { rootGraph: unknown }).rootGraph = null expect(getNodeDisplayLabel('1', 'Node #1')).toBe('Node #1') }) it('calls resolveNodeDisplayName when graph is available', () => { const mockGraph = {} const mockNode = { id: 1 } ;(app as { rootGraph: unknown }).rootGraph = mockGraph mockGetNodeByExecutionId.mockReturnValue(mockNode) mockResolveNodeDisplayName.mockReturnValue('My Checkpoint') const result = getNodeDisplayLabel('1', 'Node #1') expect(mockGetNodeByExecutionId).toHaveBeenCalledWith(mockGraph, '1') expect(result).toBe('My Checkpoint') }) }) describe('getComboValue', () => { it('returns undefined when node is not found', () => { ;(app as { rootGraph: unknown }).rootGraph = {} mockGetNodeByExecutionId.mockReturnValue(null) const result = getComboValue(makeCandidate()) expect(result).toBeUndefined() }) it('returns undefined when widget is not found', () => { ;(app as { rootGraph: unknown }).rootGraph = {} mockGetNodeByExecutionId.mockReturnValue({ widgets: [{ name: 'other_widget', value: 'test' }] }) const result = getComboValue(makeCandidate()) expect(result).toBeUndefined() }) it('returns string value directly', () => { ;(app as { rootGraph: unknown }).rootGraph = {} mockGetNodeByExecutionId.mockReturnValue({ widgets: [{ name: 'ckpt_name', value: 'v1-5.safetensors' }] }) expect(getComboValue(makeCandidate())).toBe('v1-5.safetensors') }) it('returns stringified number value', () => { ;(app as { rootGraph: unknown }).rootGraph = {} mockGetNodeByExecutionId.mockReturnValue({ widgets: [{ name: 'ckpt_name', value: 42 }] }) expect(getComboValue(makeCandidate())).toBe('42') }) it('returns undefined for unexpected types', () => { ;(app as { rootGraph: unknown }).rootGraph = {} mockGetNodeByExecutionId.mockReturnValue({ widgets: [{ name: 'ckpt_name', value: { complex: true } }] }) expect(getComboValue(makeCandidate())).toBeUndefined() }) it('returns undefined when nodeId is null', () => { const result = getComboValue(makeCandidate({ nodeId: undefined })) expect(result).toBeUndefined() }) }) describe('toggleModelExpand / isModelExpanded', () => { it('starts collapsed by default', () => { const { isModelExpanded } = useMissingModelInteractions() expect(isModelExpanded('key1')).toBe(false) }) it('toggles to expanded', () => { const { toggleModelExpand, isModelExpanded } = useMissingModelInteractions() toggleModelExpand('key1') expect(isModelExpanded('key1')).toBe(true) }) it('toggles back to collapsed', () => { const { toggleModelExpand, isModelExpanded } = useMissingModelInteractions() toggleModelExpand('key1') toggleModelExpand('key1') expect(isModelExpanded('key1')).toBe(false) }) }) describe('handleComboSelect', () => { it('sets selectedLibraryModel in store', () => { const store = useMissingModelStore() const { handleComboSelect } = useMissingModelInteractions() handleComboSelect('key1', 'model_v2.safetensors') expect(store.selectedLibraryModel['key1']).toBe('model_v2.safetensors') }) it('does not set value when undefined', () => { const store = useMissingModelStore() const { handleComboSelect } = useMissingModelInteractions() handleComboSelect('key1', undefined) expect(store.selectedLibraryModel['key1']).toBeUndefined() }) }) describe('isSelectionConfirmable', () => { it('returns false when no selection exists', () => { const { isSelectionConfirmable } = useMissingModelInteractions() expect(isSelectionConfirmable('key1')).toBe(false) }) it('returns false when download is running', () => { const store = useMissingModelStore() store.selectedLibraryModel['key1'] = 'model.safetensors' store.importTaskIds['key1'] = 'task-123' mockDownloadList.mockReturnValue([ { taskId: 'task-123', status: 'running' } ]) const { isSelectionConfirmable } = useMissingModelInteractions() expect(isSelectionConfirmable('key1')).toBe(false) }) it('returns false when importCategoryMismatch exists', () => { const store = useMissingModelStore() store.selectedLibraryModel['key1'] = 'model.safetensors' store.importCategoryMismatch['key1'] = 'loras' const { isSelectionConfirmable } = useMissingModelInteractions() expect(isSelectionConfirmable('key1')).toBe(false) }) it('returns true when selection is ready with no active download', () => { const store = useMissingModelStore() store.selectedLibraryModel['key1'] = 'model.safetensors' mockDownloadList.mockReturnValue([]) const { isSelectionConfirmable } = useMissingModelInteractions() expect(isSelectionConfirmable('key1')).toBe(true) }) }) describe('cancelLibrarySelect', () => { it('clears selectedLibraryModel and importCategoryMismatch', () => { const store = useMissingModelStore() store.selectedLibraryModel['key1'] = 'model.safetensors' store.importCategoryMismatch['key1'] = 'loras' const { cancelLibrarySelect } = useMissingModelInteractions() cancelLibrarySelect('key1') expect(store.selectedLibraryModel['key1']).toBeUndefined() expect(store.importCategoryMismatch['key1']).toBeUndefined() }) }) describe('confirmLibrarySelect', () => { it('updates widget values on referencing nodes and removes missing model', () => { const mockGraph = {} ;(app as { rootGraph: unknown }).rootGraph = mockGraph const widget1 = { name: 'ckpt_name', value: 'old_model.safetensors' } const widget2 = { name: 'ckpt_name', value: 'old_model.safetensors' } const node1 = { widgets: [widget1] } const node2 = { widgets: [widget2] } mockGetNodeByExecutionId.mockImplementation( (_graph: unknown, id: string) => { if (id === '10') return node1 if (id === '20') return node2 return null } ) const store = useMissingModelStore() store.selectedLibraryModel['key1'] = 'new_model.safetensors' store.setMissingModels([ makeCandidate({ name: 'old_model.safetensors', nodeId: '10' }), makeCandidate({ name: 'old_model.safetensors', nodeId: '20' }) ]) const removeSpy = vi.spyOn(store, 'removeMissingModelByNameOnNodes') const { confirmLibrarySelect } = useMissingModelInteractions() confirmLibrarySelect( 'key1', 'old_model.safetensors', [ { nodeId: '10', widgetName: 'ckpt_name' }, { nodeId: '20', widgetName: 'ckpt_name' } ], null ) expect(widget1.value).toBe('new_model.safetensors') expect(widget2.value).toBe('new_model.safetensors') expect(removeSpy).toHaveBeenCalledWith( 'old_model.safetensors', new Set(['10', '20']) ) expect(store.selectedLibraryModel['key1']).toBeUndefined() }) it('does nothing when no selection exists', () => { ;(app as { rootGraph: unknown }).rootGraph = {} const store = useMissingModelStore() const removeSpy = vi.spyOn(store, 'removeMissingModelByNameOnNodes') const { confirmLibrarySelect } = useMissingModelInteractions() confirmLibrarySelect('key1', 'model.safetensors', [], null) expect(removeSpy).not.toHaveBeenCalled() }) it('does nothing when graph is null', () => { ;(app as { rootGraph: unknown }).rootGraph = null const store = useMissingModelStore() store.selectedLibraryModel['key1'] = 'new.safetensors' const removeSpy = vi.spyOn(store, 'removeMissingModelByNameOnNodes') const { confirmLibrarySelect } = useMissingModelInteractions() confirmLibrarySelect('key1', 'model.safetensors', [], null) expect(removeSpy).not.toHaveBeenCalled() }) it('refreshes model cache when directory is provided', () => { ;(app as { rootGraph: unknown }).rootGraph = {} mockGetNodeByExecutionId.mockReturnValue(null) mockGetAllNodeProviders.mockReturnValue([ { nodeDef: { name: 'CheckpointLoaderSimple' } } ]) const store = useMissingModelStore() store.selectedLibraryModel['key1'] = 'new.safetensors' const { confirmLibrarySelect } = useMissingModelInteractions() confirmLibrarySelect('key1', 'model.safetensors', [], 'checkpoints') expect(mockGetAllNodeProviders).toHaveBeenCalledWith('checkpoints') }) }) describe('handleUrlInput', () => { it('clears previous state on new input', () => { const store = useMissingModelStore() store.urlMetadata['key1'] = { name: 'old' } as never store.urlErrors['key1'] = 'old error' store.urlFetching['key1'] = true const { handleUrlInput } = useMissingModelInteractions() handleUrlInput('key1', 'https://civitai.com/models/123') expect(store.urlInputs['key1']).toBe('https://civitai.com/models/123') expect(store.urlMetadata['key1']).toBeUndefined() expect(store.urlErrors['key1']).toBeUndefined() expect(store.urlFetching['key1']).toBe(false) }) it('does not set debounce timer for empty input', () => { const store = useMissingModelStore() const setTimerSpy = vi.spyOn(store, 'setDebounceTimer') const { handleUrlInput } = useMissingModelInteractions() handleUrlInput('key1', ' ') expect(setTimerSpy).not.toHaveBeenCalled() }) it('sets debounce timer for non-empty input', () => { const store = useMissingModelStore() const setTimerSpy = vi.spyOn(store, 'setDebounceTimer') const { handleUrlInput } = useMissingModelInteractions() handleUrlInput('key1', 'https://civitai.com/models/123') expect(setTimerSpy).toHaveBeenCalledWith( 'key1', expect.any(Function), 800 ) }) it('clears previous debounce timer', () => { const store = useMissingModelStore() const clearTimerSpy = vi.spyOn(store, 'clearDebounceTimer') const { handleUrlInput } = useMissingModelInteractions() handleUrlInput('key1', 'https://civitai.com/models/123') expect(clearTimerSpy).toHaveBeenCalledWith('key1') }) }) describe('getTypeMismatch', () => { it('returns null when groupDirectory is null', () => { const { getTypeMismatch } = useMissingModelInteractions() expect(getTypeMismatch('key1', null)).toBeNull() }) it('returns null when no metadata exists', () => { const { getTypeMismatch } = useMissingModelInteractions() expect(getTypeMismatch('key1', 'checkpoints')).toBeNull() }) it('returns null when metadata has no tags', () => { const store = useMissingModelStore() store.urlMetadata['key1'] = { name: 'model', tags: [] } as never const { getTypeMismatch } = useMissingModelInteractions() expect(getTypeMismatch('key1', 'checkpoints')).toBeNull() }) it('returns null when detected type matches directory', () => { const store = useMissingModelStore() store.urlMetadata['key1'] = { name: 'model', tags: ['checkpoints'] } as never const { getTypeMismatch } = useMissingModelInteractions() expect(getTypeMismatch('key1', 'checkpoints')).toBeNull() }) it('returns detected type when it differs from directory', () => { const store = useMissingModelStore() store.urlMetadata['key1'] = { name: 'model', tags: ['loras'] } as never const { getTypeMismatch } = useMissingModelInteractions() expect(getTypeMismatch('key1', 'checkpoints')).toBe('loras') }) it('returns null when tags contain no recognized model type', () => { const store = useMissingModelStore() store.urlMetadata['key1'] = { name: 'model', tags: ['other', 'random'] } as never const { getTypeMismatch } = useMissingModelInteractions() expect(getTypeMismatch('key1', 'checkpoints')).toBeNull() }) }) })