mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-03-23 22:22:11 -03:00
feat: improve file size extraction logic with primary model preference
Refactor `_extract_size_bytes` method to prioritize primary model files when calculating size. The new implementation: - Extracts size parsing into separate `parse_size` function - Adds logic to prefer files marked as both "model" type and "primary" - Falls back to first valid size if no primary model file found - Adds comprehensive tests for primary preference and fallback behavior This ensures more accurate size reporting for model files, particularly when multiple file types are present in the response.
This commit is contained in:
@@ -74,6 +74,58 @@ def make_record(*versions, should_ignore_model=False):
|
||||
)
|
||||
|
||||
|
||||
def test_extract_size_bytes_prefers_primary_model_file(tmp_path):
|
||||
db_path = tmp_path / "updates.sqlite"
|
||||
service = ModelUpdateService(str(db_path))
|
||||
|
||||
response = {
|
||||
"modelVersions": [
|
||||
{
|
||||
"id": 42,
|
||||
"files": [
|
||||
{"sizeKB": 2018.0400390625, "type": "Training Data", "primary": False},
|
||||
{
|
||||
"sizeKB": 1152322.3515625,
|
||||
"type": "Model",
|
||||
"primary": "True",
|
||||
},
|
||||
],
|
||||
"images": [],
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
versions = service._extract_versions(response)
|
||||
assert versions is not None
|
||||
assert versions[0].size_bytes == int(1152322.3515625 * 1024)
|
||||
|
||||
|
||||
def test_extract_size_bytes_falls_back_without_primary(tmp_path):
|
||||
db_path = tmp_path / "updates.sqlite"
|
||||
service = ModelUpdateService(str(db_path))
|
||||
|
||||
response = {
|
||||
"modelVersions": [
|
||||
{
|
||||
"id": 43,
|
||||
"files": [
|
||||
{
|
||||
"sizeKB": 2048,
|
||||
"type": "Training Data",
|
||||
"primary": True,
|
||||
},
|
||||
{"sizeKB": 1024, "type": "Archive", "primary": False},
|
||||
],
|
||||
"images": [],
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
versions = service._extract_versions(response)
|
||||
assert versions is not None
|
||||
assert versions[0].size_bytes == int(2048 * 1024)
|
||||
|
||||
|
||||
def test_has_update_requires_newer_version_than_library():
|
||||
record = make_record(
|
||||
make_version(5, in_library=True),
|
||||
|
||||
Reference in New Issue
Block a user