mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-08-20 12:31:27 -03:00
feat(download): per-file download status and multi-file selection (#1058)
This commit is contained in:
Will Miao
@@ -0,0 +1,317 @@
|
||||
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
|
||||
|
||||
const {
|
||||
DOWNLOAD_MANAGER_MODULE,
|
||||
MODAL_MANAGER_MODULE,
|
||||
UI_HELPERS_MODULE,
|
||||
STATE_MODULE,
|
||||
LOADING_MANAGER_MODULE,
|
||||
API_FACTORY_MODULE,
|
||||
STORAGE_HELPERS_MODULE,
|
||||
FOLDER_TREE_MANAGER_MODULE,
|
||||
I18N_HELPERS_MODULE,
|
||||
SUMMARY_MODULE,
|
||||
mockApiClient,
|
||||
mockLoadingManager,
|
||||
showToastMock,
|
||||
showDownloadBatchSummaryMock,
|
||||
resetAndReloadMock,
|
||||
} = vi.hoisted(() => {
|
||||
// Shared API client returned by the mocked getModelApiClient factory.
|
||||
const mockApiClient = {
|
||||
modelType: 'loras',
|
||||
apiConfig: {
|
||||
config: {
|
||||
displayName: 'LoRA',
|
||||
singularName: 'lora',
|
||||
},
|
||||
},
|
||||
fetchCivitaiVersions: vi.fn(),
|
||||
fetchModelRoots: vi.fn(async () => ({ roots: ['/models/loras'] })),
|
||||
fetchUnifiedFolderTree: vi.fn(async () => ({ success: false })),
|
||||
downloadModel: vi.fn(),
|
||||
cancelDownload: vi.fn(),
|
||||
getPageState: vi.fn(() => ({})),
|
||||
};
|
||||
|
||||
const mockLoadingManager = {
|
||||
showSimpleLoading: vi.fn(),
|
||||
setStatus: vi.fn(),
|
||||
hide: vi.fn(),
|
||||
restoreProgressBar: vi.fn(),
|
||||
showDownloadProgress: vi.fn(() => vi.fn()),
|
||||
showCancelButton: vi.fn(),
|
||||
};
|
||||
|
||||
return {
|
||||
DOWNLOAD_MANAGER_MODULE: new URL('../../../static/js/managers/DownloadManager.js', import.meta.url).pathname,
|
||||
MODAL_MANAGER_MODULE: new URL('../../../static/js/managers/ModalManager.js', import.meta.url).pathname,
|
||||
UI_HELPERS_MODULE: new URL('../../../static/js/utils/uiHelpers.js', import.meta.url).pathname,
|
||||
STATE_MODULE: new URL('../../../static/js/state/index.js', import.meta.url).pathname,
|
||||
LOADING_MANAGER_MODULE: new URL('../../../static/js/managers/LoadingManager.js', import.meta.url).pathname,
|
||||
API_FACTORY_MODULE: new URL('../../../static/js/api/modelApiFactory.js', import.meta.url).pathname,
|
||||
STORAGE_HELPERS_MODULE: new URL('../../../static/js/utils/storageHelpers.js', import.meta.url).pathname,
|
||||
FOLDER_TREE_MANAGER_MODULE: new URL('../../../static/js/components/FolderTreeManager.js', import.meta.url).pathname,
|
||||
I18N_HELPERS_MODULE: new URL('../../../static/js/utils/i18nHelpers.js', import.meta.url).pathname,
|
||||
SUMMARY_MODULE: new URL('../../../static/js/components/DownloadBatchSummaryModal.js', import.meta.url).pathname,
|
||||
mockApiClient,
|
||||
mockLoadingManager,
|
||||
showToastMock: vi.fn(),
|
||||
showDownloadBatchSummaryMock: vi.fn(),
|
||||
resetAndReloadMock: vi.fn(),
|
||||
};
|
||||
});
|
||||
|
||||
vi.mock(MODAL_MANAGER_MODULE, () => ({
|
||||
modalManager: {
|
||||
showModal: vi.fn(),
|
||||
closeModal: vi.fn(),
|
||||
},
|
||||
}));
|
||||
|
||||
vi.mock(UI_HELPERS_MODULE, () => ({
|
||||
showToast: showToastMock,
|
||||
setupAutoNewlineOnPaste: vi.fn(),
|
||||
}));
|
||||
|
||||
vi.mock(STATE_MODULE, () => ({
|
||||
state: {
|
||||
global: {
|
||||
settings: {},
|
||||
},
|
||||
loadingManager: mockLoadingManager,
|
||||
},
|
||||
}));
|
||||
|
||||
vi.mock(LOADING_MANAGER_MODULE, () => ({
|
||||
LoadingManager: vi.fn(() => mockLoadingManager),
|
||||
}));
|
||||
|
||||
vi.mock(API_FACTORY_MODULE, () => ({
|
||||
getModelApiClient: vi.fn(() => mockApiClient),
|
||||
resetAndReload: resetAndReloadMock,
|
||||
}));
|
||||
|
||||
vi.mock(STORAGE_HELPERS_MODULE, () => ({
|
||||
getStorageItem: vi.fn((_key, defaultValue) => defaultValue),
|
||||
setStorageItem: vi.fn(),
|
||||
}));
|
||||
|
||||
vi.mock(FOLDER_TREE_MANAGER_MODULE, () => ({
|
||||
FolderTreeManager: vi.fn(() => ({
|
||||
clearSelection: vi.fn(),
|
||||
init: vi.fn(),
|
||||
getSelectedPath: vi.fn(() => ''),
|
||||
})),
|
||||
}));
|
||||
|
||||
vi.mock(I18N_HELPERS_MODULE, () => ({
|
||||
translate: vi.fn((_, __, fallback) => fallback ?? ''),
|
||||
}));
|
||||
|
||||
vi.mock(SUMMARY_MODULE, () => ({
|
||||
showDownloadBatchSummary: showDownloadBatchSummaryMock,
|
||||
}));
|
||||
|
||||
/** DOM covering the file-selection, version and location steps. */
|
||||
function setupDownloadDom() {
|
||||
document.body.innerHTML = `
|
||||
<div id="downloadModal">
|
||||
<div class="download-step" id="urlStep"></div>
|
||||
<div class="download-step" id="versionStep"></div>
|
||||
<div class="download-step" id="fileSelectionStep"></div>
|
||||
<div class="download-step" id="downloadLocationStep"></div>
|
||||
<div id="fileSelectionList"></div>
|
||||
<div id="fileSelectionVersionName"></div>
|
||||
<button id="nextFromVersion"></button>
|
||||
<div id="downloadModalTitle"></div>
|
||||
<select id="modelRoot"></select>
|
||||
<input id="folderPath" />
|
||||
<div id="targetPathDisplay"></div>
|
||||
<input id="useDefaultPath" type="checkbox" />
|
||||
<div id="manualPathSelection"></div>
|
||||
</div>
|
||||
`;
|
||||
}
|
||||
|
||||
function makeMultiFileVersion(overrides = {}) {
|
||||
return {
|
||||
id: 201,
|
||||
name: 'Multi-file version',
|
||||
baseModel: 'SDXL',
|
||||
images: [],
|
||||
files: [
|
||||
{ id: 1001, type: 'Model', sizeKB: 2048, name: 'file-a.safetensors' },
|
||||
{ id: 1002, type: 'Model', sizeKB: 2048, name: 'file-b.safetensors' },
|
||||
{ id: 1003, type: 'Model', sizeKB: 2048, name: 'file-c.safetensors' },
|
||||
],
|
||||
createdAt: '2026-01-01T00:00:00Z',
|
||||
existsLocally: true,
|
||||
...overrides,
|
||||
};
|
||||
}
|
||||
|
||||
function getFileOption(fileId) {
|
||||
return document.querySelector(`.file-option[data-file-id="${fileId}"]`);
|
||||
}
|
||||
|
||||
describe('DownloadManager multi-select file dialog (#1058)', () => {
|
||||
let DownloadManager;
|
||||
|
||||
beforeEach(async () => {
|
||||
vi.resetModules();
|
||||
vi.clearAllMocks();
|
||||
setupDownloadDom();
|
||||
({ DownloadManager } = await import(DOWNLOAD_MANAGER_MODULE));
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
document.body.innerHTML = '';
|
||||
});
|
||||
|
||||
it('renders downloaded files disabled with an In Library tag', () => {
|
||||
const manager = new DownloadManager();
|
||||
manager.versions = [makeMultiFileVersion({
|
||||
downloadedFiles: [
|
||||
{ fileId: 1001, fileName: 'file-a.safetensors', filePath: '/models/loras/file-a.safetensors' },
|
||||
],
|
||||
})];
|
||||
|
||||
manager.showFileSelectionStep('201');
|
||||
|
||||
const downloadedOption = getFileOption('1001');
|
||||
expect(downloadedOption.classList.contains('disabled')).toBe(true);
|
||||
expect(downloadedOption.querySelector('input[type="checkbox"]').disabled).toBe(true);
|
||||
expect(downloadedOption.querySelector('.file-tag.in-library').textContent).toBe('In Library');
|
||||
|
||||
// Remaining files stay selectable
|
||||
const otherOption = getFileOption('1002');
|
||||
expect(otherOption.classList.contains('disabled')).toBe(false);
|
||||
expect(otherOption.querySelector('input[type="checkbox"]').disabled).toBe(false);
|
||||
});
|
||||
|
||||
it('ignores clicks on already-downloaded options', () => {
|
||||
const manager = new DownloadManager();
|
||||
manager.versions = [makeMultiFileVersion({
|
||||
downloadedFiles: [
|
||||
{ fileId: 1001, fileName: 'file-a.safetensors', filePath: '/models/loras/file-a.safetensors' },
|
||||
],
|
||||
})];
|
||||
|
||||
manager.showFileSelectionStep('201');
|
||||
|
||||
getFileOption('1001').click();
|
||||
|
||||
expect(getFileOption('1001').querySelector('input[type="checkbox"]').checked).toBe(false);
|
||||
expect(manager.selectedFiles).toHaveLength(0);
|
||||
});
|
||||
|
||||
it('confirmFileSelection collects multiple checked files into selectedFiles', () => {
|
||||
const manager = new DownloadManager();
|
||||
manager.apiClient = mockApiClient;
|
||||
manager.versions = [makeMultiFileVersion({ downloadedFiles: [] })];
|
||||
|
||||
manager.showFileSelectionStep('201');
|
||||
|
||||
getFileOption('1001').click();
|
||||
getFileOption('1003').click();
|
||||
|
||||
manager.confirmFileSelection();
|
||||
|
||||
expect(manager.selectedFiles.map(f => f.id)).toEqual([1001, 1003]);
|
||||
// selectedFile stays the first selected file for single-file flows
|
||||
expect(manager.selectedFile?.id).toBe(1001);
|
||||
expect(document.getElementById('fileSelectionStep').style.display).toBe('none');
|
||||
expect(document.getElementById('downloadLocationStep').style.display).toBe('block');
|
||||
});
|
||||
|
||||
it('confirmFileSelection warns when nothing is selected', () => {
|
||||
const manager = new DownloadManager();
|
||||
manager.versions = [makeMultiFileVersion({ downloadedFiles: [] })];
|
||||
|
||||
manager.showFileSelectionStep('201');
|
||||
manager.confirmFileSelection();
|
||||
|
||||
expect(showToastMock).toHaveBeenCalledWith('toast.loras.pleaseSelectFile', {}, 'error');
|
||||
expect(manager.selectedFiles).toHaveLength(0);
|
||||
expect(document.getElementById('downloadLocationStep').style.display).not.toBe('block');
|
||||
});
|
||||
|
||||
it('disables the other routing group once a file is checked and re-enables when unchecked', () => {
|
||||
const manager = new DownloadManager();
|
||||
manager.versions = [makeMultiFileVersion({
|
||||
downloadedFiles: [],
|
||||
files: [
|
||||
{ id: 1001, type: 'UNet', sizeKB: 2048, name: 'unet-a.safetensors' },
|
||||
{ id: 1002, type: 'Model', sizeKB: 2048, name: 'file-b.safetensors' },
|
||||
{ id: 1003, type: 'Model', sizeKB: 2048, name: 'file-c.safetensors' },
|
||||
],
|
||||
})];
|
||||
|
||||
manager.showFileSelectionStep('201');
|
||||
|
||||
// Checking a regular Model file disables the UNet option
|
||||
getFileOption('1002').click();
|
||||
expect(getFileOption('1001').classList.contains('group-disabled')).toBe(true);
|
||||
expect(getFileOption('1001').querySelector('input[type="checkbox"]').disabled).toBe(true);
|
||||
expect(getFileOption('1003').classList.contains('group-disabled')).toBe(false);
|
||||
|
||||
// Clicking a group-disabled option does nothing
|
||||
getFileOption('1001').click();
|
||||
expect(manager.selectedFiles.map(f => f.id)).toEqual([1002]);
|
||||
|
||||
// Unchecking everything re-enables the other group
|
||||
getFileOption('1002').click();
|
||||
expect(manager.selectedFiles).toHaveLength(0);
|
||||
expect(getFileOption('1001').classList.contains('group-disabled')).toBe(false);
|
||||
expect(getFileOption('1001').querySelector('input[type="checkbox"]').disabled).toBe(false);
|
||||
});
|
||||
|
||||
it('keeps Next enabled for a partially downloaded multi-file version', () => {
|
||||
const manager = new DownloadManager();
|
||||
manager.currentVersion = makeMultiFileVersion({
|
||||
downloadedFiles: [
|
||||
{ fileId: 1001, fileName: 'file-a.safetensors', filePath: '/models/loras/file-a.safetensors' },
|
||||
],
|
||||
});
|
||||
|
||||
manager.updateNextButtonState();
|
||||
|
||||
const nextButton = document.getElementById('nextFromVersion');
|
||||
expect(nextButton.disabled).toBe(false);
|
||||
expect(nextButton.classList.contains('disabled')).toBe(false);
|
||||
});
|
||||
|
||||
it('disables Next when every weight file is already downloaded', () => {
|
||||
const manager = new DownloadManager();
|
||||
manager.currentVersion = makeMultiFileVersion({
|
||||
downloadedFiles: [
|
||||
{ fileId: 1001, fileName: 'file-a.safetensors', filePath: '/models/loras/file-a.safetensors' },
|
||||
{ fileId: 1002, fileName: 'file-b.safetensors', filePath: '/models/loras/file-b.safetensors' },
|
||||
{ fileId: 1003, fileName: 'file-c.safetensors', filePath: '/models/loras/file-c.safetensors' },
|
||||
],
|
||||
});
|
||||
|
||||
manager.updateNextButtonState();
|
||||
|
||||
const nextButton = document.getElementById('nextFromVersion');
|
||||
expect(nextButton.disabled).toBe(true);
|
||||
expect(nextButton.classList.contains('disabled')).toBe(true);
|
||||
});
|
||||
|
||||
it('disables Next for an in-library single-file version', () => {
|
||||
const manager = new DownloadManager();
|
||||
manager.currentVersion = {
|
||||
id: 202,
|
||||
name: 'Single-file version',
|
||||
files: [{ id: 1004, type: 'Model', sizeKB: 2048, name: 'file-d.safetensors' }],
|
||||
existsLocally: true,
|
||||
downloadedFiles: [],
|
||||
};
|
||||
|
||||
manager.updateNextButtonState();
|
||||
|
||||
const nextButton = document.getElementById('nextFromVersion');
|
||||
expect(nextButton.disabled).toBe(true);
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,77 @@
|
||||
"""Unit tests for per-file downloaded-state matching (#1058)."""
|
||||
|
||||
from py.routes.handlers.model_handlers import ModelCivitaiHandler
|
||||
|
||||
|
||||
VERSION = {
|
||||
"id": 42,
|
||||
"files": [
|
||||
{
|
||||
"id": 1001,
|
||||
"name": "file-a.safetensors",
|
||||
"hashes": {"SHA256": "AAA111"},
|
||||
},
|
||||
{
|
||||
"id": 1002,
|
||||
"name": "file-b.safetensors",
|
||||
"hashes": {"SHA256": "BBB222"},
|
||||
},
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
def _entry(file_name: str, sha256: str = "", file_path: str | None = None):
|
||||
return {
|
||||
"file_name": file_name,
|
||||
"file_path": file_path or f"/models/{file_name}.safetensors",
|
||||
"sha256": sha256,
|
||||
}
|
||||
|
||||
|
||||
def test_matches_by_sha256():
|
||||
entries = [_entry("renamed-locally", "bbb222")]
|
||||
result = ModelCivitaiHandler._match_downloaded_files(VERSION, entries)
|
||||
assert result == [
|
||||
{
|
||||
"fileId": 1002,
|
||||
"fileName": "file-b.safetensors",
|
||||
"filePath": "/models/renamed-locally.safetensors",
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
def test_falls_back_to_name_when_hash_missing():
|
||||
entries = [_entry("file-a", "")]
|
||||
result = ModelCivitaiHandler._match_downloaded_files(VERSION, entries)
|
||||
assert [r["fileId"] for r in result] == [1001]
|
||||
|
||||
|
||||
def test_hash_takes_precedence_over_name():
|
||||
# Hash points at file-b while the name points at file-a: hash wins.
|
||||
entries = [_entry("file-a", "bbb222")]
|
||||
result = ModelCivitaiHandler._match_downloaded_files(VERSION, entries)
|
||||
assert [r["fileId"] for r in result] == [1002]
|
||||
|
||||
|
||||
def test_unmatched_entries_are_skipped():
|
||||
entries = [
|
||||
_entry("unrelated", "ccc333"),
|
||||
_entry("file-b", ""), # name match
|
||||
]
|
||||
result = ModelCivitaiHandler._match_downloaded_files(VERSION, entries)
|
||||
assert [r["fileId"] for r in result] == [1002]
|
||||
|
||||
|
||||
def test_multiple_files_of_same_version():
|
||||
entries = [
|
||||
_entry("file-a", "aaa111"),
|
||||
_entry("file-b", "bbb222"),
|
||||
]
|
||||
result = ModelCivitaiHandler._match_downloaded_files(VERSION, entries)
|
||||
assert [r["fileId"] for r in result] == [1001, 1002]
|
||||
|
||||
|
||||
def test_empty_inputs():
|
||||
assert ModelCivitaiHandler._match_downloaded_files(VERSION, []) == []
|
||||
assert ModelCivitaiHandler._match_downloaded_files({"id": 1}, [_entry("x")]) == []
|
||||
assert ModelCivitaiHandler._match_downloaded_files(VERSION, None) == []
|
||||
@@ -68,3 +68,94 @@ async def test_download_history_bulk_lookup(tmp_path: Path) -> None:
|
||||
5: {501, 502},
|
||||
6: {601},
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_per_file_history_tracking(tmp_path: Path) -> None:
|
||||
"""Per-file records coexist with the version-level row (#1058)."""
|
||||
db_path = tmp_path / "download-history.sqlite"
|
||||
service = DownloadedVersionHistoryService(
|
||||
str(db_path),
|
||||
settings_manager=DummySettings(),
|
||||
)
|
||||
|
||||
await service.mark_downloaded(
|
||||
"lora", 101, model_id=11, source="download",
|
||||
file_path="/models/a.safetensors", file_id=1001, file_name="a.safetensors",
|
||||
)
|
||||
await service.mark_downloaded(
|
||||
"lora", 101, model_id=11, source="download",
|
||||
file_path="/models/b.safetensors", file_id=1002, file_name="b.safetensors",
|
||||
)
|
||||
|
||||
assert await service.get_downloaded_file_ids("lora", 101) == [1001, 1002]
|
||||
# Version-level tracking remains single-row per version
|
||||
assert await service.get_downloaded_version_ids("lora", 11) == [101]
|
||||
|
||||
# Re-downloading the same file updates in place, no duplicate
|
||||
await service.mark_downloaded(
|
||||
"lora", 101, source="download", file_id=1001, file_name="a.safetensors",
|
||||
)
|
||||
assert await service.get_downloaded_file_ids("lora", 101) == [1001, 1002]
|
||||
|
||||
# Single-file deletion keeps the sibling record
|
||||
await service.mark_file_deleted("lora", 101, 1001)
|
||||
assert await service.get_downloaded_file_ids("lora", 101) == [1002]
|
||||
|
||||
# Whole-version deletion clears per-file records
|
||||
await service.mark_as_deleted("lora", 101)
|
||||
assert await service.get_downloaded_file_ids("lora", 101) == []
|
||||
assert await service.has_been_downloaded("lora", 101) is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_per_file_history_ignores_invalid_ids(tmp_path: Path) -> None:
|
||||
service = DownloadedVersionHistoryService(
|
||||
str(tmp_path / "download-history.sqlite"),
|
||||
settings_manager=DummySettings(),
|
||||
)
|
||||
|
||||
# mark_downloaded without a file id only touches the version-level table
|
||||
await service.mark_downloaded("lora", 201, model_id=21, source="scan")
|
||||
assert await service.get_downloaded_file_ids("lora", 201) == []
|
||||
|
||||
# Invalid inputs are no-ops
|
||||
await service.mark_file_deleted("lora", 201, None) # type: ignore[arg-type]
|
||||
assert await service.get_downloaded_file_ids("unknown-type", 201) == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_file_history_table_created_for_legacy_db(tmp_path: Path) -> None:
|
||||
"""Existing databases gain the per-file table via CREATE IF NOT EXISTS."""
|
||||
import sqlite3
|
||||
|
||||
db_path = tmp_path / "download-history.sqlite"
|
||||
conn = sqlite3.connect(db_path)
|
||||
conn.executescript(
|
||||
"""
|
||||
CREATE TABLE downloaded_model_versions (
|
||||
model_type TEXT NOT NULL,
|
||||
version_id INTEGER NOT NULL,
|
||||
model_id INTEGER,
|
||||
first_seen_at REAL NOT NULL,
|
||||
last_seen_at REAL NOT NULL,
|
||||
source TEXT NOT NULL,
|
||||
last_file_path TEXT,
|
||||
last_library_name TEXT,
|
||||
is_deleted_override INTEGER NOT NULL DEFAULT 0,
|
||||
PRIMARY KEY (model_type, version_id)
|
||||
);
|
||||
"""
|
||||
)
|
||||
conn.close()
|
||||
|
||||
service = DownloadedVersionHistoryService(
|
||||
str(db_path),
|
||||
settings_manager=DummySettings(),
|
||||
)
|
||||
await service.mark_downloaded(
|
||||
"lora", 301, model_id=31, source="download",
|
||||
file_id=9001, file_name="file.safetensors",
|
||||
)
|
||||
assert await service.get_downloaded_file_ids("lora", 301) == [9001]
|
||||
assert await service.has_been_downloaded("lora", 301) is True
|
||||
|
||||
@@ -61,3 +61,87 @@ async def test_model_cache_tracks_versions_by_model_id():
|
||||
assert cache.get_versions_by_model_id(2) == [
|
||||
{'versionId': 201, 'name': 'Gamma', 'fileName': 'model-b'},
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_version_files_index_tracks_multiple_files_per_version():
|
||||
"""Two downloaded files of the same version both stay indexed (#1058)."""
|
||||
item_a = {
|
||||
'file_path': '/models/v1-a.safetensors',
|
||||
'file_name': 'model-v1-a',
|
||||
'folder': '',
|
||||
'civitai': {'id': 301, 'modelId': 3, 'name': 'Multi'},
|
||||
}
|
||||
item_b = {
|
||||
'file_path': '/models/v1-b.safetensors',
|
||||
'file_name': 'model-v1-b',
|
||||
'folder': '',
|
||||
'civitai': {'id': 301, 'modelId': 3, 'name': 'Multi'},
|
||||
}
|
||||
|
||||
cache = ModelCache(
|
||||
raw_data=[item_a, item_b],
|
||||
folders=[],
|
||||
name_display_mode='model_name',
|
||||
)
|
||||
|
||||
files = cache.get_files_by_version_id(301)
|
||||
assert {f['file_path'] for f in files} == {
|
||||
'/models/v1-a.safetensors',
|
||||
'/models/v1-b.safetensors',
|
||||
}
|
||||
|
||||
# Re-adding an existing entry must not duplicate it
|
||||
cache.add_to_version_index(item_a)
|
||||
assert len(cache.get_files_by_version_id(301)) == 2
|
||||
|
||||
# Removing the indexed file re-points version_index to the sibling
|
||||
indexed = cache.version_index[301]
|
||||
sibling = item_b if indexed is item_a else item_a
|
||||
cache.remove_from_version_index(indexed)
|
||||
|
||||
assert 301 in cache.version_index
|
||||
assert cache.version_index[301]['file_path'] == sibling['file_path']
|
||||
assert cache.get_versions_by_model_id(3) == [
|
||||
{'versionId': 301, 'name': 'Multi', 'fileName': sibling['file_name']},
|
||||
]
|
||||
remaining = cache.get_files_by_version_id(301)
|
||||
assert [f['file_path'] for f in remaining] == [sibling['file_path']]
|
||||
|
||||
# Removing the last file drops the version from all indexes
|
||||
cache.remove_from_version_index(sibling)
|
||||
assert 301 not in cache.version_index
|
||||
assert cache.get_files_by_version_id(301) == []
|
||||
assert cache.get_versions_by_model_id(3) == []
|
||||
assert 3 not in cache.model_id_index
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_version_files_index_rebuild_from_raw_data():
|
||||
"""rebuild_version_index reconstructs the multi-valued index (#1058)."""
|
||||
item_a = {
|
||||
'file_path': '/models/v1-a.safetensors',
|
||||
'file_name': 'model-v1-a',
|
||||
'folder': '',
|
||||
'civitai': {'id': 401, 'modelId': 4, 'name': 'Multi'},
|
||||
}
|
||||
item_b = {
|
||||
'file_path': '/models/v1-b.safetensors',
|
||||
'file_name': 'model-v1-b',
|
||||
'folder': '',
|
||||
'civitai': {'id': 401, 'modelId': 4, 'name': 'Multi'},
|
||||
}
|
||||
|
||||
cache = ModelCache(
|
||||
raw_data=[item_a, item_b],
|
||||
folders=[],
|
||||
name_display_mode='model_name',
|
||||
)
|
||||
|
||||
cache.version_files_index = {}
|
||||
cache.rebuild_version_index()
|
||||
|
||||
assert len(cache.get_files_by_version_id(401)) == 2
|
||||
# Invalid ids normalize to empty results
|
||||
assert cache.get_files_by_version_id('not-an-int') == []
|
||||
assert cache.get_files_by_version_id(None) == []
|
||||
|
||||
Reference in New Issue
Block a user