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:
Will Miao
2025-10-29 09:16:29 +08:00
parent 569c829709
commit cf4d654c4b
2 changed files with 81 additions and 8 deletions

View File

@@ -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),