import { beforeEach, describe, expect, it, vi } from 'vitest' import { downloadModel, fetchModelMetadata, isModelDownloadable, toBrowsableUrl } from './missingModelDownload' const { fetchMock, mockIsDesktop, mockSidebarTabStore, mockStartDownload } = vi.hoisted(() => ({ fetchMock: vi.fn(), mockIsDesktop: { value: false }, mockSidebarTabStore: { activeSidebarTabId: null as string | null }, mockStartDownload: vi.fn() })) vi.stubGlobal('fetch', fetchMock) vi.mock('@/platform/distribution/types', () => ({ get isDesktop() { return mockIsDesktop.value } })) vi.mock('@/stores/electronDownloadStore', () => ({ useElectronDownloadStore: () => ({ start: mockStartDownload }) })) vi.mock('@/stores/workspace/sidebarTabStore', () => ({ useSidebarTabStore: () => mockSidebarTabStore })) let testId = 0 beforeEach(() => { vi.restoreAllMocks() vi.resetAllMocks() delete window.__comfyDesktop2 delete window.__comfyDesktop2Remote }) describe('fetchModelMetadata', () => { beforeEach(() => { mockIsDesktop.value = false mockSidebarTabStore.activeSidebarTabId = null testId++ }) it('fetches file size via HEAD for non-Civitai URLs', async () => { fetchMock.mockResolvedValueOnce({ ok: true, headers: new Headers({ 'content-length': '1048576' }) }) const url = `https://huggingface.co/org/model/resolve/main/head-${testId}.safetensors` const metadata = await fetchModelMetadata(url) expect(metadata.fileSize).toBe(1048576) expect(metadata.gatedRepoUrl).toBeNull() expect(fetchMock).toHaveBeenCalledWith(url, { method: 'HEAD' }) }) it('uses Civitai API for Civitai model URLs', async () => { const url = `https://civitai.com/api/download/models/${testId}` fetchMock.mockResolvedValueOnce({ ok: true, json: async () => ({ files: [{ sizeKB: 1024, downloadUrl: url }] }) }) const metadata = await fetchModelMetadata(url) expect(metadata.fileSize).toBe(1024 * 1024) expect(metadata.gatedRepoUrl).toBeNull() expect(fetchMock).toHaveBeenCalledWith( `https://civitai.com/api/v1/model-versions/${testId}` ) }) it('returns null fileSize when Civitai API fails', async () => { fetchMock.mockResolvedValueOnce({ ok: false }) const metadata = await fetchModelMetadata( `https://civitai.com/api/download/models/${testId}` ) expect(metadata.fileSize).toBeNull() expect(metadata.gatedRepoUrl).toBeNull() expect(fetchMock).toHaveBeenCalledTimes(1) }) it('returns gatedRepoUrl for gated HuggingFace HEAD requests (403)', async () => { fetchMock.mockResolvedValueOnce({ ok: false, status: 403 }) const metadata = await fetchModelMetadata( `https://huggingface.co/bfl/FLUX.1/resolve/main/gated-${testId}.safetensors` ) expect(metadata.gatedRepoUrl).toBe('https://huggingface.co/bfl/FLUX.1') expect(metadata.fileSize).toBeNull() }) it('does not treat HuggingFace 404/500 as gated', async () => { fetchMock.mockResolvedValueOnce({ ok: false, status: 404 }) const metadata = await fetchModelMetadata( `https://huggingface.co/org/model/resolve/main/notfound-${testId}.safetensors` ) expect(metadata.gatedRepoUrl).toBeNull() expect(metadata.fileSize).toBeNull() }) it('returns null for unrecognized Civitai URL patterns', async () => { const url = `https://civitai.com/api/v1/models/${testId}` const metadata = await fetchModelMetadata(url) expect(metadata.fileSize).toBeNull() expect(metadata.gatedRepoUrl).toBeNull() expect(fetchMock).not.toHaveBeenCalled() }) it('returns cached metadata on second call', async () => { const url = `https://huggingface.co/org/model/resolve/main/cached-${testId}.safetensors` fetchMock.mockResolvedValueOnce({ ok: true, headers: new Headers({ 'content-length': '500' }) }) const first = await fetchModelMetadata(url) const second = await fetchModelMetadata(url) expect(first.fileSize).toBe(500) expect(second.fileSize).toBe(500) expect(fetchMock).toHaveBeenCalledTimes(1) }) it('does not cache incomplete results so retries are possible', async () => { const url = `https://example.com/retry-${testId}.safetensors` fetchMock .mockResolvedValueOnce({ ok: true, headers: new Headers({}) }) .mockResolvedValueOnce({ ok: true, headers: new Headers({ 'content-length': '1024' }) }) const first = await fetchModelMetadata(url) const second = await fetchModelMetadata(url) expect(first.fileSize).toBeNull() expect(second.fileSize).toBe(1024) expect(fetchMock).toHaveBeenCalledTimes(2) }) it('deduplicates concurrent requests for the same URL', async () => { const url = `https://huggingface.co/org/model/resolve/main/dedup-${testId}.safetensors` fetchMock.mockResolvedValueOnce({ ok: true, headers: new Headers({ 'content-length': '2048' }) }) const [first, second] = await Promise.all([ fetchModelMetadata(url), fetchModelMetadata(url) ]) expect(first.fileSize).toBe(2048) expect(second.fileSize).toBe(2048) expect(fetchMock).toHaveBeenCalledTimes(1) }) }) describe('toBrowsableUrl', () => { it('replaces /resolve/ with /blob/ in HuggingFace URLs', () => { expect( toBrowsableUrl( 'https://huggingface.co/org/model/resolve/main/file.safetensors' ) ).toBe('https://huggingface.co/org/model/blob/main/file.safetensors') }) it('returns non-HuggingFace URLs unchanged', () => { const url = 'https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth' expect(toBrowsableUrl(url)).toBe(url) }) it('preserves query params in HuggingFace URLs', () => { expect( toBrowsableUrl( 'https://huggingface.co/bfl/FLUX.1/resolve/main/model.safetensors?download=true' ) ).toBe( 'https://huggingface.co/bfl/FLUX.1/blob/main/model.safetensors?download=true' ) }) it('converts Civitai api/download URL to model page', () => { expect( toBrowsableUrl('https://civitai.com/api/download/models/12345') ).toBe('https://civitai.com/models/12345') }) it('converts Civitai api/v1 URL to model page', () => { expect(toBrowsableUrl('https://civitai.com/api/v1/models/12345')).toBe( 'https://civitai.com/models/12345' ) }) it('converts civitai.red URLs to model pages', () => { expect( toBrowsableUrl('https://civitai.red/api/download/models/12345') ).toBe('https://civitai.red/models/12345') expect(toBrowsableUrl('https://civitai.red/api/v1/models/12345')).toBe( 'https://civitai.red/models/12345' ) }) }) describe('isModelDownloadable', () => { it('allows civitai.red URLs', () => { expect( isModelDownloadable({ name: 'model.safetensors', url: 'https://civitai.red/api/download/models/12345', directory: 'checkpoints' }) ).toBe(true) }) it('rejects non-allowlisted URLs', () => { expect( isModelDownloadable({ name: 'model.safetensors', url: 'https://example.com/model.safetensors', directory: 'checkpoints' }) ).toBe(false) }) }) describe('downloadModel', () => { beforeEach(() => { mockIsDesktop.value = false mockSidebarTabStore.activeSidebarTabId = null }) it('uses the Desktop2 bridge directly instead of the browser fallback', () => { const anchorClick = vi .spyOn(HTMLAnchorElement.prototype, 'click') .mockImplementation(() => {}) const desktopDownloadModel = vi .fn< (url: string, filename: string, directory: string) => Promise >() .mockResolvedValue(true) window.__comfyDesktop2 = { downloadModel: desktopDownloadModel } downloadModel( { name: 'model.safetensors', url: 'https://huggingface.co/org/model/resolve/main/model.safetensors', directory: 'checkpoints' }, { checkpoints: ['/models/checkpoints'] } ) expect(desktopDownloadModel).toHaveBeenCalledWith( 'https://huggingface.co/org/model/resolve/main/model.safetensors', 'model.safetensors', 'checkpoints' ) expect(anchorClick).not.toHaveBeenCalled() expect(mockStartDownload).not.toHaveBeenCalled() }) it('logs Desktop2 bridge failures without falling back to browser download', async () => { const anchorClick = vi .spyOn(HTMLAnchorElement.prototype, 'click') .mockImplementation(() => {}) const consoleError = vi.spyOn(console, 'error').mockImplementation(() => {}) const bridgeError = new Error('bridge failed') const desktopDownloadModel = vi .fn< (url: string, filename: string, directory: string) => Promise >() .mockRejectedValue(bridgeError) window.__comfyDesktop2 = { downloadModel: desktopDownloadModel } downloadModel( { name: 'model.safetensors', url: 'https://huggingface.co/org/model/resolve/main/model.safetensors', directory: 'checkpoints' }, { checkpoints: ['/models/checkpoints'] } ) await vi.waitFor(() => { expect(consoleError).toHaveBeenCalledWith( 'Failed to start Desktop2 model download:', bridgeError ) }) expect(anchorClick).not.toHaveBeenCalled() expect(mockStartDownload).not.toHaveBeenCalled() }) it('logs synchronous Desktop2 bridge failures without crashing', async () => { const anchorClick = vi .spyOn(HTMLAnchorElement.prototype, 'click') .mockImplementation(() => {}) const consoleError = vi.spyOn(console, 'error').mockImplementation(() => {}) const bridgeError = new Error('bridge failed before returning a promise') const desktopDownloadModel = vi .fn< (url: string, filename: string, directory: string) => Promise >() .mockImplementation(() => { throw bridgeError }) window.__comfyDesktop2 = { downloadModel: desktopDownloadModel } downloadModel( { name: 'model.safetensors', url: 'https://huggingface.co/org/model/resolve/main/model.safetensors', directory: 'checkpoints' }, { checkpoints: ['/models/checkpoints'] } ) await vi.waitFor(() => { expect(consoleError).toHaveBeenCalledWith( 'Failed to start Desktop2 model download:', bridgeError ) }) expect(anchorClick).not.toHaveBeenCalled() expect(mockStartDownload).not.toHaveBeenCalled() }) it('keeps remote Desktop2 sessions on the browser fallback', () => { const anchorClick = vi .spyOn(HTMLAnchorElement.prototype, 'click') .mockImplementation(() => {}) const desktopDownloadModel = vi .fn< (url: string, filename: string, directory: string) => Promise >() .mockResolvedValue(true) window.__comfyDesktop2 = { downloadModel: desktopDownloadModel } window.__comfyDesktop2Remote = true downloadModel( { name: 'model.safetensors', url: 'https://huggingface.co/org/model/resolve/main/model.safetensors', directory: 'checkpoints' }, { checkpoints: ['/models/checkpoints'] } ) expect(desktopDownloadModel).not.toHaveBeenCalled() expect(anchorClick).toHaveBeenCalledTimes(1) }) it('opens the model library sidebar before starting a desktop download', () => { mockIsDesktop.value = true downloadModel( { name: 'model.safetensors', url: 'https://huggingface.co/org/model/resolve/main/model.safetensors', directory: 'checkpoints' }, { checkpoints: ['/models/checkpoints'] } ) expect(mockSidebarTabStore.activeSidebarTabId).toBe('model-library') expect(mockStartDownload).toHaveBeenCalledWith({ url: 'https://huggingface.co/org/model/resolve/main/model.safetensors', savePath: '/models/checkpoints', filename: 'model.safetensors' }) }) })