mirror of
https://github.com/Comfy-Org/ComfyUI_frontend.git
synced 2026-01-26 19:09:52 +00:00
refactor: centralize all download utils across app and apply special cloud-specific behavior (#6188)
## Summary Centralized all download functionalities across app. Then changed downloadFile on the cloud distribution to stream assets via blob fetches while desktop/local retains direct anchor downloads. This fixes issue where trying to download cross-origin resources opens them in the window, potentially losing the user's unsaved changes. ## Changes - **What**: Moved `downloadBlob` into `downloadUtil`, routed all callers (3D exporter, recording manager, node template export, workflow/palette export, Litegraph save, ~~`useDownload` consumers~~) through shared helpers, and changed `downloadFile` to `fetch` first when `isCloud` so cross-origin URLs download reliably - `useDownload` is the exception since we simply cannot do model downloads through blob (forcing user to transfer the entire model data twice is bad). Fortunately on cloud, the user doesn't need to download models locally anyway. ┆Issue is synchronized with this [Notion page](https://www.notion.so/PR-6188-refactor-centralize-all-download-utils-across-app-and-apply-special-cloud-specific-behav-2946d73d365081de9f27f0994950511d) by [Unito](https://www.unito.io)
This commit is contained in:
@@ -2,15 +2,38 @@ import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
|
||||
|
||||
import { downloadFile } from '@/base/common/downloadUtil'
|
||||
|
||||
let mockIsCloud = false
|
||||
|
||||
vi.mock('@/platform/distribution/types', () => ({
|
||||
get isCloud() {
|
||||
return mockIsCloud
|
||||
}
|
||||
}))
|
||||
|
||||
// Global stubs
|
||||
const createObjectURLSpy = vi
|
||||
.spyOn(URL, 'createObjectURL')
|
||||
.mockReturnValue('blob:mock-url')
|
||||
const revokeObjectURLSpy = vi
|
||||
.spyOn(URL, 'revokeObjectURL')
|
||||
.mockImplementation(() => {})
|
||||
|
||||
describe('downloadUtil', () => {
|
||||
let mockLink: HTMLAnchorElement
|
||||
let fetchMock: ReturnType<typeof vi.fn>
|
||||
|
||||
beforeEach(() => {
|
||||
mockIsCloud = false
|
||||
fetchMock = vi.fn()
|
||||
vi.stubGlobal('fetch', fetchMock)
|
||||
createObjectURLSpy.mockClear().mockReturnValue('blob:mock-url')
|
||||
revokeObjectURLSpy.mockClear().mockImplementation(() => {})
|
||||
// Create a mock anchor element
|
||||
mockLink = {
|
||||
href: '',
|
||||
download: '',
|
||||
click: vi.fn()
|
||||
click: vi.fn(),
|
||||
style: { display: '' }
|
||||
} as unknown as HTMLAnchorElement
|
||||
|
||||
// Spy on DOM methods
|
||||
@@ -20,7 +43,7 @@ describe('downloadUtil', () => {
|
||||
})
|
||||
|
||||
afterEach(() => {
|
||||
vi.restoreAllMocks()
|
||||
vi.unstubAllGlobals()
|
||||
})
|
||||
|
||||
describe('downloadFile', () => {
|
||||
@@ -35,6 +58,8 @@ describe('downloadUtil', () => {
|
||||
expect(document.body.appendChild).toHaveBeenCalledWith(mockLink)
|
||||
expect(mockLink.click).toHaveBeenCalled()
|
||||
expect(document.body.removeChild).toHaveBeenCalledWith(mockLink)
|
||||
expect(fetchMock).not.toHaveBeenCalled()
|
||||
expect(createObjectURLSpy).not.toHaveBeenCalled()
|
||||
})
|
||||
|
||||
it('should use custom filename when provided', () => {
|
||||
@@ -45,6 +70,8 @@ describe('downloadUtil', () => {
|
||||
|
||||
expect(mockLink.href).toBe(testUrl)
|
||||
expect(mockLink.download).toBe(customFilename)
|
||||
expect(fetchMock).not.toHaveBeenCalled()
|
||||
expect(createObjectURLSpy).not.toHaveBeenCalled()
|
||||
})
|
||||
|
||||
it('should extract filename from URL query parameters', () => {
|
||||
@@ -55,6 +82,7 @@ describe('downloadUtil', () => {
|
||||
|
||||
expect(mockLink.href).toBe(testUrl)
|
||||
expect(mockLink.download).toBe('extracted-image.jpg')
|
||||
expect(createObjectURLSpy).not.toHaveBeenCalled()
|
||||
})
|
||||
|
||||
it('should use default filename when URL has no filename parameter', () => {
|
||||
@@ -64,6 +92,7 @@ describe('downloadUtil', () => {
|
||||
|
||||
expect(mockLink.href).toBe(testUrl)
|
||||
expect(mockLink.download).toBe('download.png')
|
||||
expect(createObjectURLSpy).not.toHaveBeenCalled()
|
||||
})
|
||||
|
||||
it('should handle invalid URLs gracefully', () => {
|
||||
@@ -74,6 +103,8 @@ describe('downloadUtil', () => {
|
||||
expect(mockLink.href).toBe(invalidUrl)
|
||||
expect(mockLink.download).toBe('download.png')
|
||||
expect(mockLink.click).toHaveBeenCalled()
|
||||
expect(fetchMock).not.toHaveBeenCalled()
|
||||
expect(createObjectURLSpy).not.toHaveBeenCalled()
|
||||
})
|
||||
|
||||
it('should prefer custom filename over extracted filename', () => {
|
||||
@@ -84,6 +115,7 @@ describe('downloadUtil', () => {
|
||||
downloadFile(testUrl, customFilename)
|
||||
|
||||
expect(mockLink.download).toBe(customFilename)
|
||||
expect(createObjectURLSpy).not.toHaveBeenCalled()
|
||||
})
|
||||
|
||||
it('should handle URLs with empty filename parameter', () => {
|
||||
@@ -92,6 +124,7 @@ describe('downloadUtil', () => {
|
||||
downloadFile(testUrl)
|
||||
|
||||
expect(mockLink.download).toBe('download.png')
|
||||
expect(createObjectURLSpy).not.toHaveBeenCalled()
|
||||
})
|
||||
|
||||
it('should handle relative URLs by using window.location.origin', () => {
|
||||
@@ -101,6 +134,8 @@ describe('downloadUtil', () => {
|
||||
|
||||
expect(mockLink.href).toBe(relativeUrl)
|
||||
expect(mockLink.download).toBe('relative-image.png')
|
||||
expect(fetchMock).not.toHaveBeenCalled()
|
||||
expect(createObjectURLSpy).not.toHaveBeenCalled()
|
||||
})
|
||||
|
||||
it('should clean up DOM elements after download', () => {
|
||||
@@ -111,6 +146,54 @@ describe('downloadUtil', () => {
|
||||
// Verify the element was added and then removed
|
||||
expect(document.body.appendChild).toHaveBeenCalledWith(mockLink)
|
||||
expect(document.body.removeChild).toHaveBeenCalledWith(mockLink)
|
||||
expect(fetchMock).not.toHaveBeenCalled()
|
||||
expect(createObjectURLSpy).not.toHaveBeenCalled()
|
||||
})
|
||||
|
||||
it('streams downloads via blob when running in cloud', async () => {
|
||||
mockIsCloud = true
|
||||
const testUrl = 'https://storage.googleapis.com/bucket/file.bin'
|
||||
const blob = new Blob(['test'])
|
||||
const blobFn = vi.fn().mockResolvedValue(blob)
|
||||
fetchMock.mockResolvedValue({
|
||||
ok: true,
|
||||
status: 200,
|
||||
blob: blobFn
|
||||
} as unknown as Response)
|
||||
|
||||
downloadFile(testUrl)
|
||||
|
||||
expect(fetchMock).toHaveBeenCalledWith(testUrl)
|
||||
const fetchPromise = fetchMock.mock.results[0].value as Promise<Response>
|
||||
await fetchPromise
|
||||
const blobPromise = blobFn.mock.results[0].value as Promise<Blob>
|
||||
await blobPromise
|
||||
await Promise.resolve()
|
||||
expect(blobFn).toHaveBeenCalled()
|
||||
expect(createObjectURLSpy).toHaveBeenCalledWith(blob)
|
||||
expect(revokeObjectURLSpy).toHaveBeenCalledWith('blob:mock-url')
|
||||
expect(mockLink.click).toHaveBeenCalled()
|
||||
})
|
||||
|
||||
it('logs an error when cloud fetch fails', async () => {
|
||||
mockIsCloud = true
|
||||
const testUrl = 'https://storage.googleapis.com/bucket/missing.bin'
|
||||
const consoleSpy = vi.spyOn(console, 'error').mockImplementation(() => {})
|
||||
fetchMock.mockResolvedValue({
|
||||
ok: false,
|
||||
status: 404,
|
||||
blob: vi.fn()
|
||||
} as unknown as Response)
|
||||
|
||||
downloadFile(testUrl)
|
||||
|
||||
expect(fetchMock).toHaveBeenCalledWith(testUrl)
|
||||
const fetchPromise = fetchMock.mock.results[0].value as Promise<Response>
|
||||
await fetchPromise
|
||||
await Promise.resolve()
|
||||
expect(consoleSpy).toHaveBeenCalled()
|
||||
expect(createObjectURLSpy).not.toHaveBeenCalled()
|
||||
consoleSpy.mockRestore()
|
||||
})
|
||||
})
|
||||
})
|
||||
|
||||
Reference in New Issue
Block a user