mirror of https://github.com/Comfy-Org/ComfyUI_frontend.git synced 2026-01-26 19:09:52 +00:00

Files

Alexander Brown 41ffb7c627 feat: add polling fallback for stale asset downloads (#7926 )

## Summary

Adds a polling fallback mechanism to recover from dropped WebSocket
messages during model downloads.

## Problem

When downloading models via the asset download service, status updates
are received over WebSocket. Sometimes these messages are dropped
(network issues, reconnection, etc.), causing downloads to appear
"stuck" even when they've completed on the backend.

## Solution

Periodically poll for stale downloads using the existing REST API:

- Track `lastUpdate` timestamp on each download
- Downloads without updates for 10s are considered "stale"
- Poll stale downloads every 10s via `GET /tasks/{task_id}` to check if
the asset exists
- If the asset exists with size > 0, mark the download as completed

## Implementation

- Added `lastUpdate` field to `AssetDownload` interface
- Use VueUse's `useIntervalFn` with a `watch` to auto start/stop polling
based on active downloads
- Reuse existing `handleAssetDownload` for completion (synthetic event)
- Added 9 unit tests covering the polling behavior

## Testing

- All existing tests pass
- New tests cover:
  - Basic download tracking
  - Completion/failure handling  
  - Duplicate message prevention
  - Stale download polling
  - Polling error handling

┆Issue is synchronized with this [Notion
page](https://www.notion.so/PR-7926-feat-add-polling-fallback-for-stale-asset-downloads-2e36d73d3650810ea966f5480f08b60c)
by [Unito](https://www.unito.io)

---------

Co-authored-by: Amp <amp@ampcode.com>
Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>

2026-01-09 16:23:12 -08:00

2.5 KiB

Raw Blame History

globs

**/*.test.ts

**/*.spec.ts

Vitest Patterns

Setup

Use createTestingPinia from @pinia/testing, not createPinia:

import { createTestingPinia } from '@pinia/testing'
import { setActivePinia } from 'pinia'
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'

describe('MyStore', () => {
  beforeEach(() => {
    setActivePinia(createTestingPinia({ stubActions: false }))
    vi.useFakeTimers()
    vi.resetAllMocks()
  })

  afterEach(() => {
    vi.useRealTimers()
  })
})

Why stubActions: false? By default, testing pinia stubs all actions. Set to false when testing actual store behavior.

Mock Patterns

Reset all mocks at once

beforeEach(() => {
  vi.resetAllMocks()  // Not individual mock.mockReset() calls
})

Module mocks with vi.mock()

vi.mock('@/scripts/api', () => ({
  api: {
    addEventListener: vi.fn(),
    fetchData: vi.fn()
  }
}))

vi.mock('@/services/myService', () => ({
  myService: {
    doThing: vi.fn()
  }
}))

Configure mocks in tests

import { api } from '@/scripts/api'
import { myService } from '@/services/myService'

it('handles success', () => {
  vi.mocked(myService.doThing).mockResolvedValue({ data: 'test' })
  // ... test code
})

Testing Event Listeners

When a store registers event listeners at module load time:

function getEventHandler() {
  const call = vi.mocked(api.addEventListener).mock.calls.find(
    ([event]) => event === 'my_event'
  )
  return call?.[1] as (e: CustomEvent<MyEventType>) => void
}

function dispatch(data: MyEventType) {
  const handler = getEventHandler()
  handler(new CustomEvent('my_event', { detail: data }))
}

it('handles events', () => {
  const store = useMyStore()
  dispatch({ field: 'value' })
  expect(store.items).toHaveLength(1)
})

Testing with Fake Timers

For stores with intervals, timeouts, or polling:

beforeEach(() => {
  vi.useFakeTimers()
})

afterEach(() => {
  vi.useRealTimers()
})

it('polls after delay', async () => {
  const store = useMyStore()
  store.startPolling()

  await vi.advanceTimersByTimeAsync(30000)

  expect(mockService.fetch).toHaveBeenCalled()
})

Assertion Style

Prefer .toHaveLength() over .length.toBe():

// Good
expect(store.items).toHaveLength(1)

// Avoid
expect(store.items.length).toBe(1)

Use .toMatchObject() for partial matching:

expect(store.completedItems[0]).toMatchObject({
  id: 'task-123',
  status: 'done'
})

2.5 KiB Raw Blame History