Fix tests

src/composables/useWorkflowThumbnail.ts

@@ -79,20 +79,15 @@ export const useWorkflowThumbnail = () => {
   /**
    * Store a thumbnail for a workflow
    */
-  const storeThumbnail = (
+  const storeThumbnail = async (
     workflow: ComfyWorkflow,
     width: number = DEFAULT_THUMBNAIL_WIDTH,
     height: number = DEFAULT_THUMBNAIL_HEIGHT
   ) => {
-    captureCanvasThumbnail(width, height)
-      .then((thumbnail) => {
-        if (thumbnail) {
-          workflowThumbnails.value.set(workflow.key, thumbnail)
-        }
-      })
-      .catch((error) =>
-        console.error('Failed to capture canvas thumbnail:', error)
-      )
+    const thumbnail = await captureCanvasThumbnail(width, height)
+    if (thumbnail) {
+      workflowThumbnails.value.set(workflow.key, thumbnail)
+    }
   }
 
   /**
@@ -111,7 +106,6 @@ export const useWorkflowThumbnail = () => {
       URL.revokeObjectURL(thumbnail)
     }
     workflowThumbnails.value.delete(workflowKey)
-    console.log('removed thumbnail', workflowKey)
   }
 
   /**

src/services/workflowService.ts

@@ -293,7 +293,7 @@ export const useWorkflowService = () => {
     if (activeWorkflow) {
       activeWorkflow.changeTracker.store()
       // Capture thumbnail before loading new graph
-      workflowThumbnail.storeThumbnail(activeWorkflow)
+      void workflowThumbnail.storeThumbnail(activeWorkflow)
     }
   }
 

tests-ui/tests/composables/useWorkflowThumbnail.spec.ts

@@ -1,9 +1,48 @@
-import { describe, expect, it } from 'vitest'
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
 
 import { useWorkflowThumbnail } from '@/composables/useWorkflowThumbnail'
 import { ComfyWorkflow } from '@/stores/workflowStore'
 
+global.URL.createObjectURL = vi.fn(() => 'data:image/png;base64,test')
+global.URL.revokeObjectURL = vi.fn()
+
+const mockCanvas = document.createElement('canvas')
+mockCanvas.width = 800
+mockCanvas.height = 600
+
 describe('useWorkflowThumbnail', () => {
+  beforeEach(() => {
+    vi.clearAllMocks()
+
+    const originalCreateElement = document.createElement
+    vi.spyOn(document, 'createElement').mockImplementation((tagName) => {
+      if (tagName === 'canvas') {
+        return {
+          getContext: vi.fn(() => {
+            return {
+              drawImage: vi.fn()
+            } as unknown as CanvasRenderingContext2D
+          }),
+          toBlob: vi.fn((callback) => {
+            callback(new Blob(['test'], { type: 'image/png' }))
+          })
+        } as unknown as HTMLCanvasElement
+      }
+      return originalCreateElement.call(document, tagName)
+    })
+
+    vi.spyOn(document, 'getElementById').mockImplementation((id) => {
+      if (id === 'graph-canvas') {
+        return mockCanvas
+      }
+      return null
+    })
+  })
+
+  afterEach(() => {
+    vi.restoreAllMocks()
+  })
+
   it('should capture canvas thumbnail', async () => {
     const { captureCanvasThumbnail } = useWorkflowThumbnail()
     const thumbnail = await captureCanvasThumbnail()
@@ -16,7 +55,7 @@ describe('useWorkflowThumbnail', () => {
 
     const mockWorkflow = { key: 'test-workflow-key' } as ComfyWorkflow
 
-    storeThumbnail(mockWorkflow)
+    await storeThumbnail(mockWorkflow)
 
     const thumbnail = getThumbnail('test-workflow-key')
     expect(thumbnail).toBe('data:image/png;base64,test')
@@ -28,10 +67,15 @@ describe('useWorkflowThumbnail', () => {
 
     const mockWorkflow = { key: 'test-workflow-key' } as ComfyWorkflow
 
-    storeThumbnail(mockWorkflow)
+    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()
   })
 
@@ -42,15 +86,26 @@ describe('useWorkflowThumbnail', () => {
     const mockWorkflow1 = { key: 'workflow-1' } as ComfyWorkflow
     const mockWorkflow2 = { key: 'workflow-2' } as ComfyWorkflow
 
-    storeThumbnail(mockWorkflow1)
-    storeThumbnail(mockWorkflow2)
+    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 handle canvas capture failure', async () => {
+    // Mock getElementById to return null to simulate missing canvas
+    vi.mocked(document.getElementById).mockReturnValue(null)
+
+    const { captureCanvasThumbnail } = useWorkflowThumbnail()
+    const thumbnail = await captureCanvasThumbnail()
+
+    expect(thumbnail).toBeNull()
+  })
 })