import { createTestingPinia } from '@pinia/testing' import { describe, expect, it, vi } from 'vitest' import { nextTick } from 'vue' import { createI18n } from 'vue-i18n' import type { JobListItem, JobStatus } from '@/platform/remote/comfyui/jobs/jobTypes' import { useCommandStore } from '@/stores/commandStore' import { api } from '@/scripts/api' import { TaskItemImpl, useQueueSettingsStore, useQueueStore } from '@/stores/queueStore' import { render, screen } from '@testing-library/vue' import userEvent from '@testing-library/user-event' import ComfyQueueButton from './ComfyQueueButton.vue' vi.mock('@/platform/distribution/types', () => ({ isCloud: false })) vi.mock('@/platform/telemetry', () => ({ useTelemetry: () => null })) vi.mock('@/workbench/extensions/manager/utils/graphHasMissingNodes', () => ({ graphHasMissingNodes: () => false })) vi.mock('@/scripts/app', () => ({ app: { rootGraph: {} } })) vi.mock('@/stores/workspaceStore', () => ({ useWorkspaceStore: () => ({ shiftDown: false }) })) const BatchCountEditStub = { template: '
' } const i18n = createI18n({ legacy: false, locale: 'en', messages: { en: { menu: { run: 'Run', continueIndependentBranches: 'Continue independent branches if a node fails', disabledTooltip: 'Disabled tooltip', onChange: 'On Change', onChangeTooltip: 'On change tooltip', instant: 'Instant', instantTooltip: 'Instant tooltip', stopRunInstant: 'Stop Run (Instant)', stopRunInstantTooltip: 'Stop running', runWorkflow: 'Run workflow', runWorkflowFront: 'Run workflow front', runWorkflowDisabled: 'Run workflow disabled' } } } }) function createTask(id: string, status: JobStatus): TaskItemImpl { const job: JobListItem = { id, status, create_time: Date.now(), priority: 1 } return new TaskItemImpl(job) } const stubs = { BatchCountEdit: BatchCountEditStub, DropdownMenuRoot: { template: '
' }, DropdownMenuTrigger: { template: '
' }, DropdownMenuPortal: { template: '
' }, DropdownMenuContent: { template: '
' }, DropdownMenuItem: { template: '
' }, DropdownMenuCheckboxItem: { template: '
' }, DropdownMenuItemIndicator: { template: '' } } function renderQueueButton() { const pinia = createTestingPinia({ createSpy: vi.fn }) const user = userEvent.setup() const result = render(ComfyQueueButton, { global: { plugins: [pinia, i18n], directives: { tooltip: () => {} }, stubs } }) return { ...result, user } } describe('ComfyQueueButton', () => { it('only offers continuing independent branches when supported', async () => { renderQueueButton() expect( screen.queryByTestId('continue-independent-branches') ).not.toBeInTheDocument() api.serverFeatureFlags.value = { supports_node_failure_policy: true } await nextTick() expect( screen.getByTestId('continue-independent-branches') ).toHaveTextContent('Continue independent branches if a node fails') }) it('renders the batch count control before the run button', () => { renderQueueButton() const controls = screen.getAllByTestId(/batch-count-edit|queue-button/) expect(controls[0]).toHaveAttribute('data-testid', 'batch-count-edit') expect(controls[1]).toHaveAttribute('data-testid', 'queue-button') }) it('keeps the run instant presentation while idle even with active jobs', async () => { renderQueueButton() const queueSettingsStore = useQueueSettingsStore() const queueStore = useQueueStore() queueSettingsStore.mode = 'instant-idle' queueStore.runningTasks = [createTask('run-1', 'in_progress')] await nextTick() const queueButton = screen.getByTestId('queue-button') expect(queueButton).toHaveTextContent('Run (Instant)') expect(queueButton).toHaveAttribute('data-variant', 'primary') }) it('switches to stop presentation when instant mode is armed', async () => { renderQueueButton() const queueSettingsStore = useQueueSettingsStore() queueSettingsStore.mode = 'instant-running' await nextTick() const queueButton = screen.getByTestId('queue-button') expect(queueButton).toHaveTextContent('Stop Run (Instant)') expect(queueButton).toHaveAttribute('data-variant', 'destructive') }) it('disarms instant mode without interrupting even when jobs are active', async () => { const { user } = renderQueueButton() const queueSettingsStore = useQueueSettingsStore() const queueStore = useQueueStore() const commandStore = useCommandStore() queueSettingsStore.mode = 'instant-running' queueStore.runningTasks = [createTask('run-1', 'in_progress')] await nextTick() await user.click(screen.getByTestId('queue-button')) await nextTick() expect(queueSettingsStore.mode).toBe('instant-idle') const queueButton = screen.getByTestId('queue-button') expect(queueButton).toHaveTextContent('Run (Instant)') expect(queueButton).toHaveAttribute('data-variant', 'primary') expect(commandStore.execute).not.toHaveBeenCalled() }) it('activates instant running mode when queueing again', async () => { const { user } = renderQueueButton() const queueSettingsStore = useQueueSettingsStore() const commandStore = useCommandStore() queueSettingsStore.mode = 'instant-idle' await nextTick() await user.click(screen.getByTestId('queue-button')) await nextTick() expect(queueSettingsStore.mode).toBe('instant-running') expect(commandStore.execute).toHaveBeenCalledWith('Comfy.QueuePrompt', { metadata: { subscribe_to_run: false, trigger_source: 'button' } }) }) })