mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-09-20 18:51:26 -03:00
feat(ui): show live scan progress and ETA for cache refresh
Broadcast typed scan_progress messages over /ws/fetch-progress from the manual refresh/rebuild paths of ModelScanner and RecipeScanner, and render percent, processed/total, current file name and an EMA-smoothed ETA in the loading overlay. Hardcoded refresh strings move to i18n (common.scanProgress); WS connection failure falls back to the previous static loading behavior.
This commit is contained in:
Will Miao
@@ -0,0 +1,348 @@
|
||||
import { describe, it, beforeEach, afterEach, expect, vi } from 'vitest';
|
||||
|
||||
const {
|
||||
BASE_MODEL_API_MODULE,
|
||||
STATE_MODULE,
|
||||
UI_HELPERS_MODULE,
|
||||
I18N_MODULE,
|
||||
STORAGE_MODULE,
|
||||
API_CONFIG_MODULE,
|
||||
API_FACTORY_MODULE,
|
||||
SIDEBAR_MANAGER_MODULE,
|
||||
} = vi.hoisted(() => ({
|
||||
BASE_MODEL_API_MODULE: new URL('../../../static/js/api/baseModelApi.js', import.meta.url).pathname,
|
||||
STATE_MODULE: new URL('../../../static/js/state/index.js', import.meta.url).pathname,
|
||||
UI_HELPERS_MODULE: new URL('../../../static/js/utils/uiHelpers.js', import.meta.url).pathname,
|
||||
I18N_MODULE: new URL('../../../static/js/utils/i18nHelpers.js', import.meta.url).pathname,
|
||||
STORAGE_MODULE: new URL('../../../static/js/utils/storageHelpers.js', import.meta.url).pathname,
|
||||
API_CONFIG_MODULE: new URL('../../../static/js/api/apiConfig.js', import.meta.url).pathname,
|
||||
API_FACTORY_MODULE: new URL('../../../static/js/api/modelApiFactory.js', import.meta.url).pathname,
|
||||
SIDEBAR_MANAGER_MODULE: new URL('../../../static/js/components/SidebarManager.js', import.meta.url).pathname,
|
||||
}));
|
||||
|
||||
const showToastMock = vi.fn();
|
||||
const showMock = vi.fn();
|
||||
const showCancelButtonMock = vi.fn();
|
||||
const hideMock = vi.fn();
|
||||
const restoreProgressBarMock = vi.fn();
|
||||
const setProgressMock = vi.fn();
|
||||
const setStatusMock = vi.fn();
|
||||
const resetAndReloadMock = vi.fn();
|
||||
|
||||
vi.mock(STATE_MODULE, () => ({
|
||||
state: {
|
||||
loadingManager: {
|
||||
show: showMock,
|
||||
showCancelButton: showCancelButtonMock,
|
||||
hide: hideMock,
|
||||
restoreProgressBar: restoreProgressBarMock,
|
||||
setProgress: setProgressMock,
|
||||
setStatus: setStatusMock,
|
||||
},
|
||||
},
|
||||
getCurrentPageState: vi.fn(() => ({})),
|
||||
}));
|
||||
|
||||
vi.mock(UI_HELPERS_MODULE, () => ({
|
||||
showToast: showToastMock,
|
||||
}));
|
||||
|
||||
vi.mock(I18N_MODULE, () => ({
|
||||
translate: vi.fn((key, params, fallback) => {
|
||||
if (fallback) {
|
||||
return Object.entries(params || {}).reduce(
|
||||
(text, [name, value]) => text.replaceAll(`{${name}}`, value),
|
||||
fallback
|
||||
);
|
||||
}
|
||||
return key;
|
||||
}),
|
||||
}));
|
||||
|
||||
vi.mock(STORAGE_MODULE, () => ({
|
||||
getStorageItem: vi.fn(),
|
||||
getSessionItem: vi.fn(),
|
||||
removeSessionItem: vi.fn(),
|
||||
saveMapToStorage: vi.fn(),
|
||||
}));
|
||||
|
||||
vi.mock(API_CONFIG_MODULE, () => ({
|
||||
getCompleteApiConfig: vi.fn(() => ({
|
||||
endpoints: { scan: '/api/lm/loras/scan' },
|
||||
config: { displayName: 'LoRA', singularName: 'lora' },
|
||||
})),
|
||||
getCurrentModelType: vi.fn(() => 'loras'),
|
||||
isValidModelType: vi.fn(() => true),
|
||||
DOWNLOAD_ENDPOINTS: {},
|
||||
HF_ENDPOINTS: {},
|
||||
WS_ENDPOINTS: { fetchProgress: '/ws/fetch-progress' },
|
||||
}));
|
||||
|
||||
vi.mock(API_FACTORY_MODULE, () => ({
|
||||
resetAndReload: resetAndReloadMock,
|
||||
}));
|
||||
|
||||
vi.mock(SIDEBAR_MANAGER_MODULE, () => ({
|
||||
sidebarManager: { refresh: vi.fn() },
|
||||
}));
|
||||
|
||||
class FakeWebSocket {
|
||||
static instances = [];
|
||||
static failNextConnection = false;
|
||||
|
||||
constructor(url) {
|
||||
this.url = url;
|
||||
this.onopen = null;
|
||||
this.onerror = null;
|
||||
this.onmessage = null;
|
||||
this.close = vi.fn();
|
||||
FakeWebSocket.instances.push(this);
|
||||
const shouldFail = FakeWebSocket.failNextConnection;
|
||||
FakeWebSocket.failNextConnection = false;
|
||||
queueMicrotask(() => {
|
||||
if (shouldFail) {
|
||||
this.onerror?.(new Error('connection refused'));
|
||||
} else {
|
||||
this.onopen?.();
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
emit(data) {
|
||||
this.onmessage?.({ data: JSON.stringify(data) });
|
||||
}
|
||||
}
|
||||
|
||||
async function createClient() {
|
||||
const { BaseModelApiClient } = await import(BASE_MODEL_API_MODULE);
|
||||
class TestClient extends BaseModelApiClient {}
|
||||
return new TestClient('loras');
|
||||
}
|
||||
|
||||
async function flushMicrotasks() {
|
||||
await new Promise((resolve) => setTimeout(resolve, 0));
|
||||
}
|
||||
|
||||
describe('BaseModelApiClient.refreshModels scan progress', () => {
|
||||
beforeEach(() => {
|
||||
showToastMock.mockReset();
|
||||
showMock.mockReset();
|
||||
showCancelButtonMock.mockReset();
|
||||
hideMock.mockReset();
|
||||
restoreProgressBarMock.mockReset();
|
||||
setProgressMock.mockReset();
|
||||
setStatusMock.mockReset();
|
||||
resetAndReloadMock.mockReset();
|
||||
FakeWebSocket.instances = [];
|
||||
FakeWebSocket.failNextConnection = false;
|
||||
vi.stubGlobal('WebSocket', FakeWebSocket);
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
delete global.fetch;
|
||||
vi.unstubAllGlobals();
|
||||
vi.restoreAllMocks();
|
||||
});
|
||||
|
||||
function mockFetchPending() {
|
||||
let resolveFetch;
|
||||
global.fetch = vi.fn(() => new Promise((resolve) => { resolveFetch = resolve; }));
|
||||
return {
|
||||
resolveOk: (payload = { status: 'success' }) =>
|
||||
resolveFetch({ ok: true, json: async () => payload }),
|
||||
};
|
||||
}
|
||||
|
||||
async function startRefresh(client, fullRebuild = false) {
|
||||
const promise = client.refreshModels(fullRebuild);
|
||||
await vi.waitFor(() => {
|
||||
expect(FakeWebSocket.instances.length).toBe(1);
|
||||
});
|
||||
await flushMicrotasks();
|
||||
const socket = FakeWebSocket.instances[0];
|
||||
await vi.waitFor(() => {
|
||||
expect(socket.onmessage).toBeTruthy();
|
||||
});
|
||||
return { promise, socket };
|
||||
}
|
||||
|
||||
it('shows scan progress updates from the WebSocket channel', async () => {
|
||||
const fetchControl = mockFetchPending();
|
||||
const client = await createClient();
|
||||
const { promise, socket } = await startRefresh(client);
|
||||
|
||||
expect(socket.url).toBe(`ws://${window.location.host}/ws/fetch-progress`);
|
||||
|
||||
socket.emit({
|
||||
type: 'scan_progress',
|
||||
status: 'started',
|
||||
stage: 'scan_folders',
|
||||
model_type: 'lora',
|
||||
pageType: 'loras',
|
||||
full_rebuild: false,
|
||||
progress: 0,
|
||||
});
|
||||
socket.emit({
|
||||
type: 'scan_progress',
|
||||
status: 'processing',
|
||||
stage: 'process_models',
|
||||
model_type: 'lora',
|
||||
pageType: 'loras',
|
||||
full_rebuild: false,
|
||||
progress: 50,
|
||||
processed: 5,
|
||||
total: 10,
|
||||
current_name: 'style.safetensors',
|
||||
});
|
||||
|
||||
expect(setProgressMock).toHaveBeenCalledWith(0);
|
||||
expect(setProgressMock).toHaveBeenCalledWith(50);
|
||||
const lastStatus = setStatusMock.mock.calls.at(-1)[0];
|
||||
expect(lastStatus).toContain('(5/10)');
|
||||
expect(lastStatus).toContain('style.safetensors');
|
||||
// First ETA sample only anchors the timer
|
||||
expect(lastStatus).toContain('Estimating time...');
|
||||
|
||||
fetchControl.resolveOk();
|
||||
await promise;
|
||||
|
||||
expect(resetAndReloadMock).toHaveBeenCalledWith(true);
|
||||
expect(showToastMock).toHaveBeenCalledWith(
|
||||
'toast.api.refreshComplete',
|
||||
{ action: 'Refresh' },
|
||||
'success'
|
||||
);
|
||||
expect(socket.close).toHaveBeenCalled();
|
||||
expect(hideMock).toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('ignores messages for other types or other model types', async () => {
|
||||
const fetchControl = mockFetchPending();
|
||||
const client = await createClient();
|
||||
const { promise, socket } = await startRefresh(client);
|
||||
|
||||
socket.emit({
|
||||
type: 'scan_progress',
|
||||
status: 'processing',
|
||||
stage: 'process_models',
|
||||
model_type: 'checkpoint',
|
||||
progress: 33,
|
||||
processed: 1,
|
||||
total: 3,
|
||||
});
|
||||
socket.emit({
|
||||
type: 'example_images_progress',
|
||||
status: 'running',
|
||||
model_type: 'lora',
|
||||
progress: 66,
|
||||
processed: 2,
|
||||
total: 3,
|
||||
});
|
||||
|
||||
expect(setProgressMock).not.toHaveBeenCalled();
|
||||
expect(setStatusMock).not.toHaveBeenCalled();
|
||||
|
||||
fetchControl.resolveOk();
|
||||
await promise;
|
||||
});
|
||||
|
||||
it('falls back to plain loading when the WebSocket connection fails', async () => {
|
||||
FakeWebSocket.failNextConnection = true;
|
||||
global.fetch = vi.fn().mockResolvedValue({
|
||||
ok: true,
|
||||
json: async () => ({ status: 'success' }),
|
||||
});
|
||||
|
||||
const client = await createClient();
|
||||
await client.refreshModels(true);
|
||||
|
||||
expect(global.fetch).toHaveBeenCalled();
|
||||
const [url] = global.fetch.mock.calls[0];
|
||||
expect(url.searchParams.get('full_rebuild')).toBe('true');
|
||||
expect(showMock).toHaveBeenCalledWith('Full rebuild LoRAs...', 0);
|
||||
expect(showToastMock).toHaveBeenCalledWith(
|
||||
'toast.api.refreshComplete',
|
||||
{ action: 'Full rebuild' },
|
||||
'success'
|
||||
);
|
||||
});
|
||||
|
||||
it('computes an ETA with EMA smoothing once enough samples arrive', async () => {
|
||||
const fetchControl = mockFetchPending();
|
||||
let now = 1000;
|
||||
vi.spyOn(Date, 'now').mockImplementation(() => now);
|
||||
|
||||
const client = await createClient();
|
||||
const { promise, socket } = await startRefresh(client);
|
||||
|
||||
const emitProcessing = (processed, total) => socket.emit({
|
||||
type: 'scan_progress',
|
||||
status: 'processing',
|
||||
stage: 'process_models',
|
||||
model_type: 'lora',
|
||||
progress: Math.floor((processed / total) * 100),
|
||||
processed,
|
||||
total,
|
||||
});
|
||||
|
||||
// First sample anchors the timer
|
||||
emitProcessing(1, 10);
|
||||
expect(setStatusMock.mock.calls.at(-1)[0]).toContain('Estimating time...');
|
||||
|
||||
// 100s elapsed for 2 files -> 50s per file -> 400s remaining -> ~7 min
|
||||
now = 101000;
|
||||
emitProcessing(2, 10);
|
||||
expect(setStatusMock.mock.calls.at(-1)[0]).toContain('~7 min remaining');
|
||||
|
||||
// 110s elapsed for 4 files -> EMA = 50000*0.7 + 27500*0.3 = 43250ms/file
|
||||
// remaining 6 files -> 259.5s -> ~4 min
|
||||
now = 111000;
|
||||
emitProcessing(4, 10);
|
||||
expect(setStatusMock.mock.calls.at(-1)[0]).toContain('~4 min remaining');
|
||||
|
||||
fetchControl.resolveOk();
|
||||
await promise;
|
||||
});
|
||||
|
||||
it('shows the cancelled toast when the server reports cancellation', async () => {
|
||||
const fetchControl = mockFetchPending();
|
||||
const client = await createClient();
|
||||
const { promise } = await startRefresh(client);
|
||||
|
||||
fetchControl.resolveOk({ status: 'cancelled' });
|
||||
await promise;
|
||||
|
||||
expect(showToastMock).toHaveBeenCalledWith('toast.api.operationCancelled', {}, 'info');
|
||||
expect(resetAndReloadMock).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
|
||||
describe('createScanEtaTracker / formatScanRemainingTime', () => {
|
||||
it('estimates remaining time from EMA of per-file cost', async () => {
|
||||
const { createScanEtaTracker } = await import(BASE_MODEL_API_MODULE);
|
||||
let now = 0;
|
||||
vi.spyOn(Date, 'now').mockImplementation(() => now);
|
||||
|
||||
const tracker = createScanEtaTracker();
|
||||
expect(tracker.update(1, 10)).toBe('Estimating time...');
|
||||
|
||||
now = 60000; // 60s for 3 files -> 20s/file -> 7 * 20s = 140s -> ~2 min
|
||||
expect(tracker.update(3, 10)).toBe('~2 min remaining');
|
||||
|
||||
now = 61000; // tiny delta keeps EMA near 20s/file
|
||||
expect(tracker.update(4, 10)).toBe('~2 min remaining');
|
||||
|
||||
// Done: no ETA
|
||||
expect(tracker.update(10, 10)).toBeNull();
|
||||
expect(tracker.update(0, 0)).toBeNull();
|
||||
|
||||
vi.restoreAllMocks();
|
||||
});
|
||||
|
||||
it('formats hours and sub-minute remainders', async () => {
|
||||
const { formatScanRemainingTime } = await import(BASE_MODEL_API_MODULE);
|
||||
expect(formatScanRemainingTime(30000)).toBe('Less than a minute remaining');
|
||||
expect(formatScanRemainingTime(5 * 60000)).toBe('~5 min remaining');
|
||||
expect(formatScanRemainingTime(3600000 + 30 * 60000)).toBe('~1 hr 30 min remaining');
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,285 @@
|
||||
import { describe, it, beforeEach, afterEach, expect, vi } from 'vitest';
|
||||
|
||||
const showToastMock = vi.hoisted(() => vi.fn());
|
||||
const loadingManagerMock = vi.hoisted(() => ({
|
||||
show: vi.fn(),
|
||||
hide: vi.fn(),
|
||||
restoreProgressBar: vi.fn(),
|
||||
setProgress: vi.fn(),
|
||||
setStatus: vi.fn(),
|
||||
}));
|
||||
const virtualScrollerMock = vi.hoisted(() => ({
|
||||
refreshWithData: vi.fn(),
|
||||
}));
|
||||
const getCurrentPageStateMock = vi.hoisted(() => vi.fn());
|
||||
const etaUpdateMock = vi.hoisted(() => vi.fn(() => 'ETA soon'));
|
||||
|
||||
vi.mock('../../../static/js/components/RecipeCard.js', () => ({
|
||||
RecipeCard: vi.fn(() => ({ element: document.createElement('div') })),
|
||||
}));
|
||||
|
||||
vi.mock('../../../static/js/state/index.js', () => ({
|
||||
state: {
|
||||
loadingManager: loadingManagerMock,
|
||||
virtualScroller: virtualScrollerMock,
|
||||
},
|
||||
getCurrentPageState: getCurrentPageStateMock,
|
||||
}));
|
||||
|
||||
vi.mock('../../../static/js/utils/uiHelpers.js', () => ({
|
||||
showToast: showToastMock,
|
||||
}));
|
||||
|
||||
vi.mock('../../../static/js/utils/i18nHelpers.js', () => ({
|
||||
translate: vi.fn((key, params, fallback) => {
|
||||
if (fallback) {
|
||||
return Object.entries(params || {}).reduce(
|
||||
(text, [name, value]) => text.replaceAll(`{${name}}`, value),
|
||||
fallback
|
||||
);
|
||||
}
|
||||
return key;
|
||||
}),
|
||||
}));
|
||||
|
||||
vi.mock('../../../static/js/utils/infiniteScroll.js', () => ({
|
||||
captureScrollPosition: vi.fn(),
|
||||
restoreScrollPosition: vi.fn(),
|
||||
}));
|
||||
|
||||
vi.mock('../../../static/js/api/apiConfig.js', () => ({
|
||||
WS_ENDPOINTS: { fetchProgress: '/ws/fetch-progress' },
|
||||
}));
|
||||
|
||||
vi.mock('../../../static/js/utils/scanEtaUtils.js', () => ({
|
||||
createScanEtaTracker: () => ({ update: etaUpdateMock }),
|
||||
}));
|
||||
|
||||
import { refreshRecipes } from '../../../static/js/api/recipeApi.js';
|
||||
|
||||
class FakeWebSocket {
|
||||
static instances = [];
|
||||
static failNextConnection = false;
|
||||
|
||||
constructor(url) {
|
||||
this.url = url;
|
||||
this.onopen = null;
|
||||
this.onerror = null;
|
||||
this.onmessage = null;
|
||||
this.close = vi.fn();
|
||||
FakeWebSocket.instances.push(this);
|
||||
const shouldFail = FakeWebSocket.failNextConnection;
|
||||
FakeWebSocket.failNextConnection = false;
|
||||
queueMicrotask(() => {
|
||||
if (shouldFail) {
|
||||
this.onerror?.(new Error('connection refused'));
|
||||
} else {
|
||||
this.onopen?.();
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
emit(data) {
|
||||
this.onmessage?.({ data: JSON.stringify(data) });
|
||||
}
|
||||
}
|
||||
|
||||
async function flushMicrotasks() {
|
||||
await new Promise((resolve) => setTimeout(resolve, 0));
|
||||
}
|
||||
|
||||
describe('refreshRecipes scan progress', () => {
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
getCurrentPageStateMock.mockReturnValue({
|
||||
pageSize: 50,
|
||||
currentPage: 1,
|
||||
hasMore: true,
|
||||
isLoading: false,
|
||||
sortBy: 'date:desc',
|
||||
showFavoritesOnly: false,
|
||||
activeFolder: null,
|
||||
searchOptions: { recursive: true },
|
||||
customFilter: { active: false },
|
||||
filters: {},
|
||||
});
|
||||
FakeWebSocket.instances = [];
|
||||
FakeWebSocket.failNextConnection = false;
|
||||
vi.stubGlobal('WebSocket', FakeWebSocket);
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
delete global.fetch;
|
||||
vi.unstubAllGlobals();
|
||||
});
|
||||
|
||||
function mockFetchPendingScan() {
|
||||
let resolveScan;
|
||||
global.fetch = vi.fn((input) => {
|
||||
const url = String(input);
|
||||
if (url.includes('/scan')) {
|
||||
return new Promise((resolve) => { resolveScan = resolve; });
|
||||
}
|
||||
// Recipe list reload after the scan completes
|
||||
return Promise.resolve({
|
||||
ok: true,
|
||||
json: async () => ({ items: [], total: 0, total_pages: 0 }),
|
||||
});
|
||||
});
|
||||
return {
|
||||
resolveOk: (payload = { status: 'success' }) =>
|
||||
resolveScan({ ok: true, json: async () => payload }),
|
||||
resolveNotOk: () =>
|
||||
resolveScan({ ok: false, status: 500, statusText: 'Server Error' }),
|
||||
};
|
||||
}
|
||||
|
||||
async function startRefresh(fullRebuild = true) {
|
||||
const promise = refreshRecipes(fullRebuild);
|
||||
await vi.waitFor(() => {
|
||||
expect(FakeWebSocket.instances.length).toBe(1);
|
||||
});
|
||||
await flushMicrotasks();
|
||||
const socket = FakeWebSocket.instances[0];
|
||||
await vi.waitFor(() => {
|
||||
expect(socket.onmessage).toBeTruthy();
|
||||
});
|
||||
return { promise, socket };
|
||||
}
|
||||
|
||||
it('shows scan progress updates from the WebSocket channel', async () => {
|
||||
const fetchControl = mockFetchPendingScan();
|
||||
const { promise, socket } = await startRefresh();
|
||||
|
||||
expect(socket.url).toBe(`ws://${window.location.host}/ws/fetch-progress`);
|
||||
|
||||
socket.emit({
|
||||
type: 'scan_progress',
|
||||
status: 'started',
|
||||
stage: 'scan_folders',
|
||||
model_type: 'recipe',
|
||||
pageType: 'recipes',
|
||||
full_rebuild: true,
|
||||
progress: 0,
|
||||
});
|
||||
socket.emit({
|
||||
type: 'scan_progress',
|
||||
status: 'processing',
|
||||
stage: 'process_models',
|
||||
model_type: 'recipe',
|
||||
pageType: 'recipes',
|
||||
full_rebuild: true,
|
||||
progress: 50,
|
||||
processed: 5,
|
||||
total: 10,
|
||||
current_name: 'style.recipe.json',
|
||||
});
|
||||
|
||||
expect(loadingManagerMock.setProgress).toHaveBeenCalledWith(0);
|
||||
expect(loadingManagerMock.setProgress).toHaveBeenCalledWith(50);
|
||||
const lastStatus = loadingManagerMock.setStatus.mock.calls.at(-1)[0];
|
||||
expect(lastStatus).toContain('(5/10)');
|
||||
expect(lastStatus).toContain('style.recipe.json');
|
||||
expect(lastStatus).toContain('ETA soon');
|
||||
expect(etaUpdateMock).toHaveBeenCalledWith(5, 10);
|
||||
|
||||
fetchControl.resolveOk();
|
||||
await promise;
|
||||
|
||||
expect(showToastMock).toHaveBeenCalledWith(
|
||||
'toast.api.refreshComplete',
|
||||
{ action: 'Full rebuild' },
|
||||
'success'
|
||||
);
|
||||
expect(socket.close).toHaveBeenCalled();
|
||||
expect(loadingManagerMock.hide).toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('ignores messages for other types or other model types', async () => {
|
||||
const fetchControl = mockFetchPendingScan();
|
||||
const { promise, socket } = await startRefresh();
|
||||
|
||||
socket.emit({
|
||||
type: 'scan_progress',
|
||||
status: 'processing',
|
||||
stage: 'process_models',
|
||||
model_type: 'lora',
|
||||
progress: 33,
|
||||
processed: 1,
|
||||
total: 3,
|
||||
});
|
||||
socket.emit({
|
||||
type: 'example_images_progress',
|
||||
status: 'running',
|
||||
model_type: 'recipe',
|
||||
progress: 66,
|
||||
processed: 2,
|
||||
total: 3,
|
||||
});
|
||||
|
||||
expect(loadingManagerMock.setProgress).not.toHaveBeenCalled();
|
||||
expect(loadingManagerMock.setStatus).not.toHaveBeenCalled();
|
||||
|
||||
fetchControl.resolveOk();
|
||||
await promise;
|
||||
});
|
||||
|
||||
it('falls back to plain loading when the WebSocket connection fails', async () => {
|
||||
FakeWebSocket.failNextConnection = true;
|
||||
global.fetch = vi.fn((input) => {
|
||||
const url = String(input);
|
||||
if (url.includes('/scan')) {
|
||||
return Promise.resolve({
|
||||
ok: true,
|
||||
json: async () => ({ status: 'success' }),
|
||||
});
|
||||
}
|
||||
return Promise.resolve({
|
||||
ok: true,
|
||||
json: async () => ({ items: [], total: 0, total_pages: 0 }),
|
||||
});
|
||||
});
|
||||
|
||||
await refreshRecipes(false);
|
||||
|
||||
expect(global.fetch).toHaveBeenCalled();
|
||||
const [url] = global.fetch.mock.calls[0];
|
||||
expect(url.searchParams.get('full_rebuild')).toBe('false');
|
||||
expect(loadingManagerMock.show).toHaveBeenCalledWith('Refreshing Recipes...', 0);
|
||||
expect(showToastMock).toHaveBeenCalledWith(
|
||||
'toast.api.refreshComplete',
|
||||
{ action: 'Refresh' },
|
||||
'success'
|
||||
);
|
||||
});
|
||||
|
||||
it('shows the cancelled toast when the server reports cancellation', async () => {
|
||||
const fetchControl = mockFetchPendingScan();
|
||||
const { promise } = await startRefresh();
|
||||
|
||||
fetchControl.resolveOk({ status: 'cancelled' });
|
||||
await promise;
|
||||
|
||||
expect(showToastMock).toHaveBeenCalledWith('toast.api.operationCancelled', {}, 'info');
|
||||
expect(showToastMock).not.toHaveBeenCalledWith(
|
||||
'toast.api.refreshComplete',
|
||||
expect.anything(),
|
||||
expect.anything()
|
||||
);
|
||||
});
|
||||
|
||||
it('reports refresh failures through the error toast', async () => {
|
||||
const fetchControl = mockFetchPendingScan();
|
||||
const { promise } = await startRefresh();
|
||||
|
||||
fetchControl.resolveNotOk();
|
||||
await promise;
|
||||
|
||||
expect(showToastMock).toHaveBeenCalledWith(
|
||||
'toast.api.refreshFailed',
|
||||
{ action: 'rebuild', type: 'recipe' },
|
||||
'error'
|
||||
);
|
||||
expect(loadingManagerMock.hide).toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
@@ -30,10 +30,14 @@ from py.utils.models import BaseModelMetadata
|
||||
class RecordingWebSocketManager:
|
||||
def __init__(self) -> None:
|
||||
self.payloads: List[Dict[str, Any]] = []
|
||||
self.broadcasts: List[Dict[str, Any]] = []
|
||||
|
||||
async def broadcast_init_progress(self, payload: Dict[str, Any]) -> None:
|
||||
self.payloads.append(payload)
|
||||
|
||||
async def broadcast(self, payload: Dict[str, Any]) -> None:
|
||||
self.broadcasts.append(payload)
|
||||
|
||||
|
||||
def _normalize_path(path: Path) -> str:
|
||||
return str(path).replace(os.sep, "/")
|
||||
@@ -1395,3 +1399,185 @@ async def test_get_all_folders_invalidated_after_move(tmp_path: Path):
|
||||
assert "new" in all_folders
|
||||
assert "new/deep" in all_folders
|
||||
assert set(cache.folders) <= set(all_folders)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_initialize_cache_broadcasts_scan_progress(tmp_path: Path, monkeypatch):
|
||||
_create_files(tmp_path)
|
||||
scanner = DummyScanner(tmp_path)
|
||||
|
||||
ws_stub = RecordingWebSocketManager()
|
||||
monkeypatch.setattr(model_scanner, "ws_manager", ws_stub)
|
||||
|
||||
await scanner._initialize_cache()
|
||||
|
||||
messages = ws_stub.broadcasts
|
||||
assert messages, "expected scan_progress broadcasts"
|
||||
|
||||
started = messages[0]
|
||||
assert started["type"] == "scan_progress"
|
||||
assert started["status"] == "started"
|
||||
assert started["stage"] == "scan_folders"
|
||||
assert started["progress"] == 0
|
||||
assert started["model_type"] == "dummy"
|
||||
assert started["pageType"] == "dummy"
|
||||
assert started["full_rebuild"] is True
|
||||
|
||||
count_messages = [m for m in messages if m["stage"] == "count_models"]
|
||||
assert count_messages and count_messages[0]["total"] == 3
|
||||
|
||||
process_messages = [
|
||||
m for m in messages
|
||||
if m["stage"] == "process_models" and m["status"] == "processing"
|
||||
]
|
||||
assert process_messages, "expected at least one process_models update"
|
||||
final_process = process_messages[-1]
|
||||
assert final_process["processed"] == 3
|
||||
assert final_process["total"] == 3
|
||||
assert final_process["current_name"].endswith(".txt")
|
||||
for message in process_messages:
|
||||
assert 0 < message["progress"] <= 99
|
||||
|
||||
stages = [m["stage"] for m in messages]
|
||||
assert "finalizing" in stages
|
||||
completed = messages[-1]
|
||||
assert completed["status"] == "completed"
|
||||
assert completed["progress"] == 100
|
||||
assert completed["elapsed_seconds"] >= 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_initialize_cache_broadcasts_cancelled(tmp_path: Path, monkeypatch):
|
||||
_create_files(tmp_path)
|
||||
scanner = DummyScanner(tmp_path)
|
||||
|
||||
ws_stub = RecordingWebSocketManager()
|
||||
monkeypatch.setattr(model_scanner, "ws_manager", ws_stub)
|
||||
|
||||
original_process = DummyScanner._process_model_file
|
||||
|
||||
async def cancelling_process(self, file_path, root_path, **kwargs):
|
||||
scanner.cancel_task()
|
||||
return await original_process(self, file_path, root_path, **kwargs)
|
||||
|
||||
monkeypatch.setattr(DummyScanner, "_process_model_file", cancelling_process)
|
||||
|
||||
await scanner._initialize_cache()
|
||||
|
||||
messages = ws_stub.broadcasts
|
||||
assert messages[0]["status"] == "started"
|
||||
assert messages[-1]["status"] == "cancelled"
|
||||
assert messages[-1]["elapsed_seconds"] >= 0
|
||||
assert not any(m["status"] == "completed" for m in messages)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_initialize_cache_broadcasts_error(tmp_path: Path, monkeypatch):
|
||||
scanner = DummyScanner(tmp_path)
|
||||
|
||||
ws_stub = RecordingWebSocketManager()
|
||||
monkeypatch.setattr(model_scanner, "ws_manager", ws_stub)
|
||||
|
||||
async def raising_gather(**_kwargs):
|
||||
raise RuntimeError("boom")
|
||||
|
||||
monkeypatch.setattr(scanner, "_gather_model_data", raising_gather)
|
||||
|
||||
await scanner._initialize_cache()
|
||||
|
||||
messages = ws_stub.broadcasts
|
||||
assert messages[0]["status"] == "started"
|
||||
assert messages[-1]["status"] == "error"
|
||||
assert messages[-1]["error"] == "boom"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reconcile_cache_broadcasts_scan_progress(tmp_path: Path, monkeypatch):
|
||||
_create_files(tmp_path)
|
||||
scanner = DummyScanner(tmp_path)
|
||||
await scanner._initialize_cache()
|
||||
|
||||
ws_stub = RecordingWebSocketManager()
|
||||
monkeypatch.setattr(model_scanner, "ws_manager", ws_stub)
|
||||
|
||||
new_file = tmp_path / "three.txt"
|
||||
new_file.write_text("three", encoding="utf-8")
|
||||
|
||||
await scanner._reconcile_cache()
|
||||
|
||||
messages = ws_stub.broadcasts
|
||||
assert messages, "expected scan_progress broadcasts"
|
||||
|
||||
started = messages[0]
|
||||
assert started["type"] == "scan_progress"
|
||||
assert started["status"] == "started"
|
||||
assert started["stage"] == "reconcile_scan"
|
||||
assert started["progress"] == 0
|
||||
assert started["full_rebuild"] is False
|
||||
|
||||
process_messages = [
|
||||
m for m in messages
|
||||
if m["stage"] == "process_new" and m["status"] == "processing"
|
||||
]
|
||||
assert process_messages, "expected process_new progress updates"
|
||||
assert process_messages[-1]["processed"] == 1
|
||||
assert process_messages[-1]["total"] == 1
|
||||
assert process_messages[-1]["current_name"] == "three.txt"
|
||||
|
||||
completed = messages[-1]
|
||||
assert completed["status"] == "completed"
|
||||
assert completed["progress"] == 100
|
||||
assert completed["added"] == 1
|
||||
assert completed["removed"] == 0
|
||||
assert completed["elapsed_seconds"] >= 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reconcile_cache_broadcasts_cancelled(tmp_path: Path, monkeypatch):
|
||||
_create_files(tmp_path)
|
||||
scanner = DummyScanner(tmp_path)
|
||||
await scanner._initialize_cache()
|
||||
|
||||
ws_stub = RecordingWebSocketManager()
|
||||
monkeypatch.setattr(model_scanner, "ws_manager", ws_stub)
|
||||
|
||||
new_file = tmp_path / "four.txt"
|
||||
new_file.write_text("four", encoding="utf-8")
|
||||
|
||||
original_process = DummyScanner._process_model_file
|
||||
|
||||
async def cancelling_process(self, file_path, root_path, **kwargs):
|
||||
scanner.cancel_task()
|
||||
return await original_process(self, file_path, root_path, **kwargs)
|
||||
|
||||
monkeypatch.setattr(DummyScanner, "_process_model_file", cancelling_process)
|
||||
|
||||
await scanner._reconcile_cache()
|
||||
|
||||
messages = ws_stub.broadcasts
|
||||
assert messages[0]["status"] == "started"
|
||||
assert messages[-1]["status"] == "cancelled"
|
||||
assert messages[-1]["elapsed_seconds"] >= 0
|
||||
assert not any(m["status"] == "completed" for m in messages)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reconcile_cache_broadcasts_error(tmp_path: Path, monkeypatch):
|
||||
_create_files(tmp_path)
|
||||
scanner = DummyScanner(tmp_path)
|
||||
await scanner._initialize_cache()
|
||||
|
||||
ws_stub = RecordingWebSocketManager()
|
||||
monkeypatch.setattr(model_scanner, "ws_manager", ws_stub)
|
||||
|
||||
def raising_walk(*_args, **_kwargs):
|
||||
raise RuntimeError("walk failed")
|
||||
|
||||
monkeypatch.setattr(model_scanner.os, "walk", raising_walk)
|
||||
|
||||
await scanner._reconcile_cache()
|
||||
|
||||
messages = ws_stub.broadcasts
|
||||
assert messages[0]["status"] == "started"
|
||||
assert messages[-1]["status"] == "error"
|
||||
assert messages[-1]["error"] == "walk failed"
|
||||
|
||||
@@ -9,6 +9,7 @@ import pytest
|
||||
|
||||
from py.config import config
|
||||
from py.services import model_scanner as model_scanner_module
|
||||
from py.services import recipe_scanner as recipe_scanner_module
|
||||
from py.services.model_cache import ModelCache
|
||||
from py.services.model_hash_index import ModelHashIndex
|
||||
from py.services.model_scanner import CacheBuildResult, ModelScanner
|
||||
@@ -4965,3 +4966,133 @@ async def test_find_all_duplicate_recipes_include_prompt_missing_gen_params(reci
|
||||
groups = await scanner.find_all_duplicate_recipes(include_prompt=True)
|
||||
# Recipes without gen_params/prompt normalize to empty prompt and match
|
||||
assert groups == {"abc:0.8\x1f": ["r1", "r2"]}
|
||||
|
||||
|
||||
class RecordingRecipeWebSocketManager:
|
||||
"""Minimal ws_manager stand-in that records broadcasts."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.payloads: list[Dict[str, Any]] = []
|
||||
self.broadcasts: list[Dict[str, Any]] = []
|
||||
|
||||
async def broadcast_init_progress(self, payload: Dict[str, Any]) -> None:
|
||||
self.payloads.append(payload)
|
||||
|
||||
async def broadcast(self, payload: Dict[str, Any]) -> None:
|
||||
self.broadcasts.append(payload)
|
||||
|
||||
|
||||
def _write_progress_recipe_files(recipes_dir: Path, count: int) -> None:
|
||||
recipes_dir.mkdir(parents=True, exist_ok=True)
|
||||
for idx in range(count):
|
||||
recipe_path = recipes_dir / f"progress-recipe-{idx}.recipe.json"
|
||||
recipe_path.write_text(
|
||||
json.dumps(
|
||||
{
|
||||
"id": f"progress-recipe-{idx}",
|
||||
"file_path": str(recipes_dir / f"img-{idx}.png"),
|
||||
"title": f"Recipe {idx}",
|
||||
"modified": 0.0,
|
||||
"created_date": 0.0,
|
||||
"loras": [],
|
||||
}
|
||||
),
|
||||
encoding="utf-8",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_force_refresh_broadcasts_scan_progress(
|
||||
tmp_path: Path, monkeypatch, recipe_scanner
|
||||
):
|
||||
scanner, _stub = recipe_scanner
|
||||
recipes_dir = Path(config.loras_roots[0]) / "recipes"
|
||||
_write_progress_recipe_files(recipes_dir, 3)
|
||||
|
||||
ws_stub = RecordingRecipeWebSocketManager()
|
||||
monkeypatch.setattr(recipe_scanner_module, "ws_manager", ws_stub)
|
||||
|
||||
await scanner.get_cached_data(force_refresh=True)
|
||||
# Wait for the FTS index build so no background task outlives the loop.
|
||||
if scanner._fts_index_task:
|
||||
await scanner._fts_index_task
|
||||
|
||||
messages = ws_stub.broadcasts
|
||||
assert messages, "expected scan_progress broadcasts"
|
||||
|
||||
started = messages[0]
|
||||
assert started["type"] == "scan_progress"
|
||||
assert started["status"] == "started"
|
||||
assert started["stage"] == "scan_folders"
|
||||
assert started["progress"] == 0
|
||||
assert started["model_type"] == "recipe"
|
||||
assert started["pageType"] == "recipes"
|
||||
assert started["full_rebuild"] is True
|
||||
|
||||
count_messages = [m for m in messages if m["stage"] == "count_models"]
|
||||
assert count_messages and count_messages[0]["total"] == 3
|
||||
|
||||
process_messages = [
|
||||
m
|
||||
for m in messages
|
||||
if m["stage"] == "process_models" and m["status"] == "processing"
|
||||
]
|
||||
assert process_messages, "expected at least one process_models update"
|
||||
final_process = process_messages[-1]
|
||||
assert final_process["processed"] == 3
|
||||
assert final_process["total"] == 3
|
||||
assert final_process["current_name"].endswith(".recipe.json")
|
||||
for message in process_messages:
|
||||
assert 0 < message["progress"] <= 99
|
||||
|
||||
completed = messages[-1]
|
||||
assert completed["status"] == "completed"
|
||||
assert completed["progress"] == 100
|
||||
assert completed["elapsed_seconds"] >= 0
|
||||
assert completed["total"] == 3
|
||||
|
||||
|
||||
def test_sync_init_without_report_progress_does_not_broadcast(
|
||||
tmp_path: Path, monkeypatch, recipe_scanner
|
||||
):
|
||||
"""Startup path (initialize_in_background) must not emit scan_progress."""
|
||||
scanner, _stub = recipe_scanner
|
||||
recipes_dir = Path(config.loras_roots[0]) / "recipes"
|
||||
_write_progress_recipe_files(recipes_dir, 2)
|
||||
|
||||
ws_stub = RecordingRecipeWebSocketManager()
|
||||
monkeypatch.setattr(recipe_scanner_module, "ws_manager", ws_stub)
|
||||
|
||||
# Invalidate the persistent cache so the sync path performs a full
|
||||
# directory scan, exactly like a force refresh but without progress
|
||||
# reporting (this is how initialize_in_background invokes it).
|
||||
scanner._persistent_cache.save_cache([], {})
|
||||
|
||||
scanner._initialize_recipe_cache_sync()
|
||||
|
||||
assert ws_stub.broadcasts == []
|
||||
|
||||
|
||||
def test_sync_init_reports_error_broadcast(
|
||||
tmp_path: Path, monkeypatch, recipe_scanner
|
||||
):
|
||||
scanner, _stub = recipe_scanner
|
||||
recipes_dir = Path(config.loras_roots[0]) / "recipes"
|
||||
_write_progress_recipe_files(recipes_dir, 1)
|
||||
|
||||
ws_stub = RecordingRecipeWebSocketManager()
|
||||
monkeypatch.setattr(recipe_scanner_module, "ws_manager", ws_stub)
|
||||
|
||||
scanner._persistent_cache.save_cache([], {})
|
||||
|
||||
def raising_scan(self, recipes_dir, progress_loop=None):
|
||||
raise RuntimeError("boom")
|
||||
|
||||
monkeypatch.setattr(RecipeScanner, "_full_directory_scan_sync", raising_scan)
|
||||
|
||||
scanner._initialize_recipe_cache_sync(report_progress=True)
|
||||
|
||||
messages = ws_stub.broadcasts
|
||||
assert messages[0]["status"] == "started"
|
||||
assert messages[-1]["status"] == "error"
|
||||
assert messages[-1]["error"] == "boom"
|
||||
|
||||
Reference in New Issue
Block a user