import { createPinia, setActivePinia } from 'pinia' import { beforeEach, describe, expect, it, vi } from 'vitest' import type { ComfyWorkflow } from '@/platform/workflow/management/stores/workflowStore' import { useWorkflowStore } from '@/platform/workflow/management/stores/workflowStore' vi.mock('@/renderer/core/thumbnail/graphThumbnailRenderer', () => ({ createGraphThumbnail: vi.fn() })) vi.mock('@/scripts/api', () => ({ api: { moveUserData: vi.fn(), listUserDataFullInfo: vi.fn(), addEventListener: vi.fn(), getUserData: vi.fn(), storeUserData: vi.fn(), apiURL: vi.fn((path: string) => `/api${path}`) } })) const { useWorkflowThumbnail } = await import('@/renderer/core/thumbnail/useWorkflowThumbnail') const { createGraphThumbnail } = await import('@/renderer/core/thumbnail/graphThumbnailRenderer') const { api } = await import('@/scripts/api') describe('useWorkflowThumbnail', () => { let workflowStore: ReturnType beforeEach(() => { setActivePinia(createPinia()) workflowStore = useWorkflowStore() // Clear any existing thumbnails from previous tests BEFORE mocking const { clearAllThumbnails } = useWorkflowThumbnail() clearAllThumbnails() // Now set up mocks vi.clearAllMocks() global.URL.createObjectURL = vi.fn(() => 'data:image/png;base64,test') global.URL.revokeObjectURL = vi.fn() // Mock API responses vi.mocked(api.moveUserData).mockResolvedValue({ status: 200 } as Response) // Default createGraphThumbnail to return test value vi.mocked(createGraphThumbnail).mockReturnValue( 'data:image/png;base64,test' ) }) it('should capture minimap thumbnail', async () => { const { createMinimapPreview } = useWorkflowThumbnail() const thumbnail = await createMinimapPreview() expect(createGraphThumbnail).toHaveBeenCalledOnce() expect(thumbnail).toBe('data:image/png;base64,test') }) it('should store and retrieve thumbnails', async () => { const { storeThumbnail, getThumbnail } = useWorkflowThumbnail() const mockWorkflow = { key: 'test-workflow-key' } as ComfyWorkflow await storeThumbnail(mockWorkflow) const thumbnail = getThumbnail('test-workflow-key') expect(thumbnail).toBe('data:image/png;base64,test') }) it('should clear thumbnail', async () => { const { storeThumbnail, getThumbnail, clearThumbnail } = useWorkflowThumbnail() const mockWorkflow = { key: 'test-workflow-key' } as ComfyWorkflow await storeThumbnail(mockWorkflow) expect(getThumbnail('test-workflow-key')).toBeDefined() clearThumbnail('test-workflow-key') expect(URL.revokeObjectURL).toHaveBeenCalledWith( 'data:image/png;base64,test' ) expect(getThumbnail('test-workflow-key')).toBeUndefined() }) it('should clear all thumbnails', async () => { const { storeThumbnail, getThumbnail, clearAllThumbnails } = useWorkflowThumbnail() const mockWorkflow1 = { key: 'workflow-1' } as ComfyWorkflow const mockWorkflow2 = { key: 'workflow-2' } as ComfyWorkflow await storeThumbnail(mockWorkflow1) await storeThumbnail(mockWorkflow2) expect(getThumbnail('workflow-1')).toBeDefined() expect(getThumbnail('workflow-2')).toBeDefined() clearAllThumbnails() expect(URL.revokeObjectURL).toHaveBeenCalledTimes(2) expect(getThumbnail('workflow-1')).toBeUndefined() expect(getThumbnail('workflow-2')).toBeUndefined() }) it('should automatically handle thumbnail cleanup when workflow is renamed', async () => { const { storeThumbnail, getThumbnail, workflowThumbnails } = useWorkflowThumbnail() // Create a temporary workflow const workflow = workflowStore.createTemporary('test-workflow.json') const originalKey = workflow.key // Store thumbnail for the workflow await storeThumbnail(workflow) expect(getThumbnail(originalKey)).toBe('data:image/png;base64,test') expect(workflowThumbnails.value.size).toBe(1) // Rename the workflow - this should automatically handle thumbnail cleanup const newPath = 'workflows/renamed-workflow.json' await workflowStore.renameWorkflow(workflow, newPath) const newKey = workflow.key // The workflow's key should now be the new path // The thumbnail should be moved from old key to new key expect(getThumbnail(originalKey)).toBeUndefined() expect(getThumbnail(newKey)).toBe('data:image/png;base64,test') expect(workflowThumbnails.value.size).toBe(1) // No URL should be revoked since we're moving the thumbnail, not deleting it expect(URL.revokeObjectURL).not.toHaveBeenCalled() }) it('should properly revoke old URL when storing thumbnail over existing one', async () => { const { storeThumbnail, getThumbnail } = useWorkflowThumbnail() const mockWorkflow = { key: 'test-workflow' } as ComfyWorkflow // Store first thumbnail await storeThumbnail(mockWorkflow) const firstThumbnail = getThumbnail('test-workflow') expect(firstThumbnail).toBe('data:image/png;base64,test') // Reset the mock to track new calls and create different URL vi.clearAllMocks() global.URL.createObjectURL = vi.fn(() => 'data:image/png;base64,test2') vi.mocked(createGraphThumbnail).mockReturnValue( 'data:image/png;base64,test2' ) // Store second thumbnail for same workflow - should revoke the first URL await storeThumbnail(mockWorkflow) const secondThumbnail = getThumbnail('test-workflow') expect(secondThumbnail).toBe('data:image/png;base64,test2') // URL.revokeObjectURL should have been called for the first thumbnail expect(URL.revokeObjectURL).toHaveBeenCalledWith( 'data:image/png;base64,test' ) expect(URL.revokeObjectURL).toHaveBeenCalledTimes(1) }) it('should clear thumbnail when workflow is deleted', async () => { const { storeThumbnail, getThumbnail, workflowThumbnails } = useWorkflowThumbnail() // Create a workflow and store thumbnail const workflow = workflowStore.createTemporary('test-delete.json') await storeThumbnail(workflow) expect(getThumbnail(workflow.key)).toBe('data:image/png;base64,test') expect(workflowThumbnails.value.size).toBe(1) // Delete the workflow - this should clear the thumbnail await workflowStore.deleteWorkflow(workflow) // Thumbnail should be cleared and URL revoked expect(getThumbnail(workflow.key)).toBeUndefined() expect(workflowThumbnails.value.size).toBe(0) expect(URL.revokeObjectURL).toHaveBeenCalledWith( 'data:image/png;base64,test' ) }) it('should clear thumbnail when temporary workflow is closed', async () => { const { storeThumbnail, getThumbnail, workflowThumbnails } = useWorkflowThumbnail() // Create a temporary workflow and store thumbnail const workflow = workflowStore.createTemporary('temp-workflow.json') await storeThumbnail(workflow) expect(getThumbnail(workflow.key)).toBe('data:image/png;base64,test') expect(workflowThumbnails.value.size).toBe(1) // Close the workflow - this should clear the thumbnail for temporary workflows await workflowStore.closeWorkflow(workflow) // Thumbnail should be cleared and URL revoked expect(getThumbnail(workflow.key)).toBeUndefined() expect(workflowThumbnails.value.size).toBe(0) expect(URL.revokeObjectURL).toHaveBeenCalledWith( 'data:image/png;base64,test' ) }) it('should handle multiple renames without leaking', async () => { const { storeThumbnail, getThumbnail, workflowThumbnails } = useWorkflowThumbnail() // Create workflow and store thumbnail const workflow = workflowStore.createTemporary('original.json') await storeThumbnail(workflow) const originalKey = workflow.key expect(getThumbnail(originalKey)).toBe('data:image/png;base64,test') expect(workflowThumbnails.value.size).toBe(1) // Rename multiple times await workflowStore.renameWorkflow(workflow, 'workflows/renamed1.json') const firstRenameKey = workflow.key expect(getThumbnail(originalKey)).toBeUndefined() expect(getThumbnail(firstRenameKey)).toBe('data:image/png;base64,test') expect(workflowThumbnails.value.size).toBe(1) await workflowStore.renameWorkflow(workflow, 'workflows/renamed2.json') const secondRenameKey = workflow.key expect(getThumbnail(originalKey)).toBeUndefined() expect(getThumbnail(firstRenameKey)).toBeUndefined() expect(getThumbnail(secondRenameKey)).toBe('data:image/png;base64,test') expect(workflowThumbnails.value.size).toBe(1) // No URLs should be revoked since we're just moving thumbnails expect(URL.revokeObjectURL).not.toHaveBeenCalled() }) it('should handle edge cases like empty keys or invalid operations', async () => { const { getThumbnail, clearThumbnail, moveWorkflowThumbnail, workflowThumbnails } = useWorkflowThumbnail() // Test getting non-existent thumbnail expect(getThumbnail('non-existent')).toBeUndefined() // Test clearing non-existent thumbnail (should not throw) expect(() => clearThumbnail('non-existent')).not.toThrow() expect(URL.revokeObjectURL).not.toHaveBeenCalled() // Test moving non-existent thumbnail (should not throw) expect(() => moveWorkflowThumbnail('non-existent', 'target')).not.toThrow() expect(workflowThumbnails.value.size).toBe(0) // Test moving to same key (should not cause issues) const { storeThumbnail } = useWorkflowThumbnail() const mockWorkflow = { key: 'test-key' } as ComfyWorkflow await storeThumbnail(mockWorkflow) expect(workflowThumbnails.value.size).toBe(1) moveWorkflowThumbnail('test-key', 'test-key') expect(workflowThumbnails.value.size).toBe(1) expect(getThumbnail('test-key')).toBe('data:image/png;base64,test') }) })