ComfyUI_frontend/tests-ui/tests/store/comfyManagerStore.test.ts

import { createPinia, setActivePinia } from 'pinia'
import { beforeEach, describe, expect, it, vi } from 'vitest'
import { nextTick, ref } from 'vue'

import { useComfyManagerService } from '@/services/comfyManagerService'
import { useComfyManagerStore } from '@/stores/comfyManagerStore'
import {
  InstalledPacksResponse,
  ManagerChannel,
  ManagerDatabaseSource,
  ManagerPackInstalled
} from '@/types/comfyManagerTypes'

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

vi.mock('@/services/dialogService', () => ({
  useDialogService: () => ({
    showManagerProgressDialog: vi.fn()
  })
}))

vi.mock('@/composables/useManagerQueue', () => {
  const enqueueTaskMock = vi.fn()

  return {
    useManagerQueue: () => ({
      statusMessage: ref(''),
      allTasksDone: ref(false),
      enqueueTask: enqueueTaskMock,
      uncompletedCount: ref(0)
    }),
    enqueueTask: enqueueTaskMock
  }
})

vi.mock('@/composables/useServerLogs', () => ({
  useServerLogs: () => ({
    startListening: vi.fn(),
    stopListening: vi.fn(),
    logs: ref([])
  })
}))

vi.mock('vue-i18n', () => ({
  useI18n: () => ({
    t: vi.fn((key) => key)
  }),
  createI18n: vi.fn(() => ({
    global: {
      t: vi.fn((key) => key)
    }
  }))
}))

interface EnabledDisabledTestCase {
  desc: string
  installed: Record<string, ManagerPackInstalled>
  expectState: 'enabled' | 'disabled'
  /** @default 'name' */
  packName?: string
}

describe('useComfyManagerStore', () => {
  let mockManagerService: ReturnType<typeof useComfyManagerService>

  const triggerPacksChange = async (
    installedPacks: InstalledPacksResponse,
    store: ReturnType<typeof useComfyManagerStore>
  ) => {
    // Simulate change in value to properly trigger watchers. Required even for immediate watchers.
    store.installedPacks = {}
    await nextTick()
    store.installedPacks = installedPacks
  }

  beforeEach(() => {
    setActivePinia(createPinia())
    vi.clearAllMocks()
    mockManagerService = {
      isLoading: ref(false),
      error: ref(null),
      startQueue: vi.fn().mockResolvedValue(null),
      resetQueue: vi.fn().mockResolvedValue(null),
      getQueueStatus: vi.fn().mockResolvedValue(null),
      listInstalledPacks: vi.fn().mockResolvedValue({}),
      getImportFailInfo: vi.fn().mockResolvedValue(null),
      installPack: vi.fn().mockResolvedValue(null),
      uninstallPack: vi.fn().mockResolvedValue(null),
      enablePack: vi.fn().mockResolvedValue(null),
      disablePack: vi.fn().mockResolvedValue(null),
      updatePack: vi.fn().mockResolvedValue(null),
      updateAllPacks: vi.fn().mockResolvedValue(null),
      rebootComfyUI: vi.fn().mockResolvedValue(null),
      isLegacyManagerUI: vi.fn().mockResolvedValue(false)
    }

    vi.mocked(useComfyManagerService).mockReturnValue(mockManagerService)
  })

  const testCases: EnabledDisabledTestCase[] = [
    {
      desc: 'Two enabled versions',
      installed: {
        'name@1_0_2': {
          enabled: true,
          cnr_id: 'name',
          ver: '1.0.2',
          aux_id: undefined
        },
        name: { enabled: true, cnr_id: 'name', ver: '1.0.0', aux_id: undefined }
      },
      expectState: 'enabled'
    },
    {
      desc: 'Two disabled versions',
      installed: {
        'name@1_0_2': {
          enabled: false,
          cnr_id: 'name',
          ver: '1.0.2',
          aux_id: undefined
        },
        name: {
          enabled: false,
          cnr_id: 'name',
          ver: '1.0.0',
          aux_id: undefined
        }
      },
      expectState: 'disabled'
    },
    {
      desc: 'Enabled version and pinned disabled version',
      installed: {
        'name@1_0_2': {
          enabled: false,
          cnr_id: 'name',
          ver: '1.0.2',
          aux_id: undefined
        },
        name: { enabled: true, cnr_id: 'name', ver: '1.0.0', aux_id: undefined }
      },
      expectState: 'enabled'
    },
    {
      desc: 'Disabled version and pinned enabled version',
      installed: {
        'name@1_0_2': {
          enabled: true,
          cnr_id: 'name',
          ver: '1.0.2',
          aux_id: undefined
        },
        name: {
          enabled: false,
          cnr_id: 'name',
          ver: '1.0.0',
          aux_id: undefined
        }
      },
      expectState: 'enabled'
    },
    {
      desc: 'Pinned enabled version, Pinned disabled version',
      installed: {
        'name@1_0_2': {
          enabled: false,
          cnr_id: 'name',
          ver: '1.0.2',
          aux_id: undefined
        },
        'name@1_0_3': {
          enabled: true,
          cnr_id: 'name',
          ver: '1.0.3',
          aux_id: undefined
        }
      },
      expectState: 'enabled'
    },
    {
      desc: 'Two enabled non-CNR versions',
      packName: 'author/name',
      installed: {
        'author/name@1_0_2': {
          enabled: true,
          aux_id: 'author/name',
          cnr_id: undefined,
          ver: '1.0.2'
        },
        'author/name': {
          enabled: true,
          aux_id: 'author/name',
          cnr_id: undefined,
          ver: '1.0.0'
        }
      },
      expectState: 'enabled'
    },
    {
      desc: 'Two disabled non-CNR versions',
      packName: 'author/name',
      installed: {
        'author/name@1_0_2': {
          enabled: false,
          aux_id: 'author/name',
          cnr_id: undefined,
          ver: '1.0.2'
        },
        'author/name': {
          enabled: false,
          aux_id: 'author/name',
          cnr_id: undefined,
          ver: '1.0.0'
        }
      },
      expectState: 'disabled'
    },
    {
      desc: 'Non-CNR disabled version, CNR enabled version',
      packName: 'author/name',
      installed: {
        'author/name': {
          enabled: false,
          aux_id: 'author/name',
          cnr_id: undefined,
          ver: '1.0.0'
        },
        'author/name@1_0_2': {
          enabled: true,
          cnr_id: 'author/name',
          ver: '1.0.2',
          aux_id: undefined
        }
      },
      expectState: 'enabled'
    },
    {
      desc: 'Disabled non-CNR version, CNR disabled version',
      packName: 'author/name',
      installed: {
        'author/name': {
          enabled: false,
          aux_id: 'author/name',
          cnr_id: undefined, // non-CNR pack
          ver: '1.0.0'
        },
        'author/name@1_0_2': {
          enabled: false,
          cnr_id: 'author/name', // CNR pack
          ver: '1.0.2',
          aux_id: undefined
        }
      },
      expectState: 'disabled'
    },
    {
      desc: 'Enabled non-CNR version, two versions',
      packName: 'author/name',
      installed: {
        'author/name': {
          enabled: true,
          aux_id: 'author/name',
          cnr_id: undefined,
          ver: '1.0.0'
        },
        'author/name@1_0_2': {
          enabled: true,
          cnr_id: 'author/name',
          ver: '1.0.2',
          aux_id: undefined
        }
      },
      expectState: 'enabled'
    },
    {
      desc: 'Enabled CNR version',
      packName: 'name',
      installed: {
        name: { enabled: true, cnr_id: 'name', ver: '1.0.0', aux_id: undefined }
      },
      expectState: 'enabled'
    },
    {
      desc: 'Disabled CNR version',
      packName: 'name',
      installed: {
        name: {
          enabled: false,
          cnr_id: 'name',
          ver: '1.0.0',
          aux_id: undefined
        }
      },
      expectState: 'disabled'
    },
    {
      desc: 'Enabled non-CNR version',
      packName: 'author/name',
      installed: {
        'author/name': {
          enabled: true,
          aux_id: 'author/name',
          cnr_id: undefined,
          ver: '1.0.0'
        }
      },
      expectState: 'enabled'
    },
    {
      desc: 'Disabled non-CNR version',
      packName: 'author/name',
      installed: {
        'author/name': {
          enabled: false,
          aux_id: 'author/name',
          cnr_id: undefined,
          ver: '1.0.0'
        }
      },
      expectState: 'disabled'
    },
    {
      desc: 'Pack not installed',
      installed: {
        'a different pack': {
          enabled: true,
          cnr_id: 'a different pack',
          ver: '1.0.0',
          aux_id: undefined
        }
      },
      expectState: 'disabled'
    }
  ]

  describe('isPackEnabled', () => {
    it.each(testCases)(
      '$expectState when $desc',
      async ({ installed, expectState, packName }) => {
        packName ??= 'name'

        const store = useComfyManagerStore()
        await triggerPacksChange(installed, store)

        const enabled = expectState === 'enabled'
        expect(store.isPackEnabled(packName)).toBe(enabled)
      }
    )
  })

  describe('isPackInstalling', () => {
    it('should return false for packs not being installed', () => {
      const store = useComfyManagerStore()
      expect(store.isPackInstalling('test-pack')).toBe(false)
      expect(store.isPackInstalling(undefined)).toBe(false)
      expect(store.isPackInstalling('')).toBe(false)
    })

    it('should track pack as installing when installPack is called', async () => {
      const store = useComfyManagerStore()

      // Call installPack
      await store.installPack.call({
        id: 'test-pack',
        repository: 'https://github.com/test/test-pack',
        channel: ManagerChannel.DEV,
        mode: ManagerDatabaseSource.CACHE,
        selected_version: 'latest',
        version: 'latest'
      })

      // Check that the pack is marked as installing
      expect(store.isPackInstalling('test-pack')).toBe(true)
    })

    it('should remove pack from installing list when explicitly removed', async () => {
      const store = useComfyManagerStore()

      // Call installPack
      await store.installPack.call({
        id: 'test-pack',
        repository: 'https://github.com/test/test-pack',
        channel: ManagerChannel.DEV,
        mode: ManagerDatabaseSource.CACHE,
        selected_version: 'latest',
        version: 'latest'
      })

      // Verify pack is installing
      expect(store.isPackInstalling('test-pack')).toBe(true)

      // Call installPack again for another pack to demonstrate multiple installs
      await store.installPack.call({
        id: 'another-pack',
        repository: 'https://github.com/test/another-pack',
        channel: ManagerChannel.DEV,
        mode: ManagerDatabaseSource.CACHE,
        selected_version: 'latest',
        version: 'latest'
      })

      // Both should be installing
      expect(store.isPackInstalling('test-pack')).toBe(true)
      expect(store.isPackInstalling('another-pack')).toBe(true)
    })

    it('should track multiple packs installing independently', async () => {
      const store = useComfyManagerStore()

      // Install pack 1
      await store.installPack.call({
        id: 'pack-1',
        repository: 'https://github.com/test/pack-1',
        channel: ManagerChannel.DEV,
        mode: ManagerDatabaseSource.CACHE,
        selected_version: 'latest',
        version: 'latest'
      })

      // Install pack 2
      await store.installPack.call({
        id: 'pack-2',
        repository: 'https://github.com/test/pack-2',
        channel: ManagerChannel.DEV,
        mode: ManagerDatabaseSource.CACHE,
        selected_version: 'latest',
        version: 'latest'
      })

      // Both should be installing
      expect(store.isPackInstalling('pack-1')).toBe(true)
      expect(store.isPackInstalling('pack-2')).toBe(true)
      expect(store.isPackInstalling('pack-3')).toBe(false)
    })
  })
})