diff --git a/core/services/civitai_search_service.py b/core/services/civitai_search_service.py index ad022fa..0bb6185 100644 --- a/core/services/civitai_search_service.py +++ b/core/services/civitai_search_service.py @@ -2,7 +2,7 @@ from .. import path_utils from ..routes.context import RouteContext -from ..sources.common import is_remote_link_marked_dead +from ..sources.common import collect_download_urls from .model_utils import CivitAISearchDependencies, ModelServiceDependencies @@ -361,53 +361,14 @@ class CivitAISearchService(ModelServiceDependencies): ) def collect_result_download_urls(result): - urls = [] - if is_remote_link_marked_dead(result): - return urls - expected_filename = result.get("filename") or filename - dead_urls = set() - mirrors = result.get("mirrors") or [] - if not isinstance(mirrors, list): - mirrors = [mirrors] - for mirror in mirrors: - if not isinstance(mirror, dict): - continue - if is_remote_link_marked_dead(mirror): - dead_url = str(mirror.get("url") or "").strip() - if dead_url.startswith(("http://", "https://")): - dead_urls.add(dead_url) - continue - url = str(mirror.get("url") or "").strip() - mirror_filename = ( - mirror.get("filename") - or mirror.get("name") - or expected_filename - ) - if ( - looks_like_model_file(url, mirror_filename) - and url not in urls - ): - urls.append(url) - raw_urls = result.get("download_urls") or [] - if not isinstance(raw_urls, list): - raw_urls = [raw_urls] - for raw_url in raw_urls: - url = str(raw_url or "").strip() - if ( - looks_like_model_file(url, expected_filename) - and url not in dead_urls - and url not in urls - ): - urls.append(url) - for key in ("download_url", "downloadUrl"): - url = str(result.get(key) or "").strip() - if ( - looks_like_model_file(url, expected_filename) - and url not in dead_urls - and url not in urls - ): - urls.append(url) - return urls + return collect_download_urls( + result, + model_file_predicate=looks_like_model_file, + expected_filename=filename, + include_download_urls=True, + skip_dead_item=True, + download_url_keys=("download_url", "downloadUrl"), + ) def remote_download_url_is_alive(url): headers = { diff --git a/core/sources/civarchive.py b/core/sources/civarchive.py index 2450052..c194271 100644 --- a/core/sources/civarchive.py +++ b/core/sources/civarchive.py @@ -51,6 +51,7 @@ from ..type_utils import ( from .common import ( build_custom_url_result, build_unified_search_result, + collect_download_urls, is_remote_link_marked_dead, ) @@ -676,7 +677,7 @@ def _resolve_file_size_bytes( if size: return size - urls = download_urls if download_urls is not None else _collect_download_urls_unified( + urls = download_urls if download_urls is not None else _collect_archive_download_urls( file_info, **_NORMALIZED_DOWNLOAD_URL_OPTIONS, ) @@ -874,64 +875,22 @@ def _select_primary_file(files: List[Dict[str, Any]]) -> Optional[Dict[str, Any] return select_primary_model_file(files) -def _collect_download_urls_unified( +def _collect_archive_download_urls( file_info: Dict[str, Any], prioritize_civitai_last: bool = False, skip_file_if_dead: bool = False, check_download_urls_list: bool = False, download_url_keys: Tuple[str, ...] = ("downloadUrl",), ) -> List[str]: - urls: List[str] = [] - if skip_file_if_dead and is_remote_link_marked_dead(file_info): - return urls - - expected_filename = file_info.get("filename") or file_info.get("name") or "" - dead_urls = set() - mirrors = file_info.get("mirrors") or [] - if not isinstance(mirrors, list): - mirrors = [mirrors] - - for mirror in mirrors: - if not isinstance(mirror, dict): - continue - if is_remote_link_marked_dead(mirror): - dead_url = _normalize_download_url(mirror.get("url")) - if dead_url: - dead_urls.add(dead_url) - continue - url = _normalize_download_url(mirror.get("url")) - mirror_filename = mirror.get("filename") or mirror.get("name") or expected_filename - if ( - url - and _download_url_looks_like_model_file(url, mirror_filename) - and url not in urls - ): - urls.append(url) - - if check_download_urls_list: - raw_urls = file_info.get("download_urls") or [] - if not isinstance(raw_urls, list): - raw_urls = [raw_urls] - for raw_url in raw_urls: - url = _normalize_download_url(raw_url) - if ( - url - and url not in dead_urls - and _download_url_looks_like_model_file(url, expected_filename) - and url not in urls - ): - urls.append(url) - - if not is_remote_link_marked_dead(file_info): - for key in download_url_keys: - url = _normalize_download_url(file_info.get(key)) - if ( - url - and url not in dead_urls - and _download_url_looks_like_model_file(url, expected_filename) - and url not in urls - ): - urls.append(url) + urls = collect_download_urls( + file_info, + model_file_predicate=_download_url_looks_like_model_file, + expected_filename=file_info.get("filename") or file_info.get("name") or "", + url_normalizer=_normalize_download_url, + include_download_urls=check_download_urls_list, + skip_dead_item=skip_file_if_dead, + download_url_keys=download_url_keys, + ) if prioritize_civitai_last: non_civitai_urls = [ @@ -1010,7 +969,7 @@ def _normalize_archive_mirrors(file_info: Dict[str, Any]) -> List[Dict[str, Any] def _normalize_archive_file(file_info: Dict[str, Any], model_id: Optional[int], version_id: Optional[int]) -> Dict[str, Any]: transformed = _transform_file_entry(file_info) - download_urls = _collect_download_urls_unified( + download_urls = _collect_archive_download_urls( transformed, **_ARCHIVE_DOWNLOAD_URL_OPTIONS, ) @@ -1211,7 +1170,7 @@ def _select_primary_model_file(files: List[Dict[str, Any]]) -> Optional[Dict[str file_info for file_info in files if isinstance(file_info, dict) - and _collect_download_urls_unified( + and _collect_archive_download_urls( file_info, **_NORMALIZED_DOWNLOAD_URL_OPTIONS, ) @@ -1225,7 +1184,7 @@ def _build_result_from_normalized_version( file_info: Dict[str, Any], match_type: str, ) -> Optional[Dict[str, Any]]: - download_urls = _collect_download_urls_unified( + download_urls = _collect_archive_download_urls( file_info, **_NORMALIZED_DOWNLOAD_URL_OPTIONS, ) @@ -1464,7 +1423,7 @@ def _select_hash_page_file( continue if is_remote_link_marked_dead(file_info): continue - if not _collect_download_urls_unified( + if not _collect_archive_download_urls( file_info, **_NORMALIZED_DOWNLOAD_URL_OPTIONS, ): @@ -1500,7 +1459,7 @@ def _select_model_details_file( for file_info in files: if not isinstance(file_info, dict): continue - if not _collect_download_urls_unified( + if not _collect_archive_download_urls( file_info, **_NORMALIZED_DOWNLOAD_URL_OPTIONS, ): @@ -1561,7 +1520,7 @@ def _prefer_query_matching_mirror( download_urls = [] if preferred_url: download_urls.append(preferred_url) - for url in _collect_download_urls_unified( + for url in _collect_archive_download_urls( file_info, **_NORMALIZED_DOWNLOAD_URL_OPTIONS, ): @@ -1657,7 +1616,7 @@ def _build_result_from_payload( if exact_only and best_confidence < 100.0: return None - download_urls = _collect_download_urls_unified( + download_urls = _collect_archive_download_urls( selected_file, **_ARCHIVE_DOWNLOAD_URL_OPTIONS, ) diff --git a/core/sources/common.py b/core/sources/common.py index 9aa529a..681f516 100644 --- a/core/sources/common.py +++ b/core/sources/common.py @@ -2,7 +2,7 @@ Common utilities and result builder functions for external model metadata sources. """ -from typing import Any, Dict, List, Optional +from typing import Any, Callable, Dict, List, Optional, Tuple from ..type_utils import build_search_result, normalize_sha256 @@ -24,6 +24,77 @@ def is_remote_link_marked_dead(item: Any) -> bool: ) +def collect_download_urls( + item: Dict[str, Any], + *, + model_file_predicate: Callable[[str, str], bool], + expected_filename: str = "", + url_normalizer: Optional[Callable[[Any], Optional[str]]] = None, + include_download_urls: bool = False, + skip_dead_item: bool = False, + download_url_keys: Tuple[str, ...] = ("downloadUrl",), +) -> List[str]: + """Collect valid, ordered, and deduplicated model download URLs.""" + if not isinstance(item, dict): + return [] + if skip_dead_item and is_remote_link_marked_dead(item): + return [] + + def normalize_url(value: Any) -> Optional[str]: + normalized = ( + url_normalizer(value) + if url_normalizer is not None + else str(value or "").strip() + ) + return normalized or None + + expected = item.get("filename") or expected_filename + urls: List[str] = [] + dead_urls = set() + mirrors = item.get("mirrors") or [] + if not isinstance(mirrors, list): + mirrors = [mirrors] + + for mirror in mirrors: + if not isinstance(mirror, dict): + continue + url = normalize_url(mirror.get("url")) + if is_remote_link_marked_dead(mirror): + if url and url.startswith(("http://", "https://")): + dead_urls.add(url) + continue + mirror_filename = mirror.get("filename") or mirror.get("name") or expected + if url and model_file_predicate(url, mirror_filename) and url not in urls: + urls.append(url) + + if include_download_urls: + raw_urls = item.get("download_urls") or [] + if not isinstance(raw_urls, list): + raw_urls = [raw_urls] + for raw_url in raw_urls: + url = normalize_url(raw_url) + if ( + url + and url not in dead_urls + and model_file_predicate(url, expected) + and url not in urls + ): + urls.append(url) + + if not is_remote_link_marked_dead(item): + for key in download_url_keys: + url = normalize_url(item.get(key)) + if ( + url + and url not in dead_urls + and model_file_predicate(url, expected) + and url not in urls + ): + urls.append(url) + + return urls + + def normalize_hashes_dict(hashes: Optional[Dict[str, Any]]) -> Dict[str, str]: """ Normalize a hashes dictionary (e.g. from CivitAI API or CivArchive) so that diff --git a/tests/test_refactoring_targets.py b/tests/test_refactoring_targets.py index e709fce..d811b7a 100644 --- a/tests/test_refactoring_targets.py +++ b/tests/test_refactoring_targets.py @@ -2,6 +2,7 @@ import importlib import json import os import sys +import tempfile import unittest from typing import ClassVar from unittest.mock import AsyncMock, MagicMock, patch @@ -1290,6 +1291,78 @@ class TestRefactoringTargets(unittest.IsolatedAsyncioTestCase): "resolved_path is outside configured model directories", ) + async def test_civitai_service_filters_dead_and_duplicate_download_urls(self): + expected_hash = "a" * 64 + dead_url = "https://dead.example/model.safetensors" + mirror_url = "https://mirror.example/model.safetensors" + listed_url = "https://listed.example/model.safetensors" + direct_url = "https://direct.example/model.safetensors" + result = { + "source": "civarchive", + "from_metadata": True, + "filename": "model.safetensors", + "sha256": expected_hash, + "mirrors": [ + {"url": dead_url, "isDead": True}, + {"url": mirror_url}, + ], + "download_urls": [dead_url, listed_url, mirror_url], + "download_url": direct_url, + } + + def to_bool(value, default=False): + return default if value is None else bool(value) + + response = MagicMock(status_code=200) + request_public_url = MagicMock(return_value=(response, mirror_url, {})) + resolve_civarchive_by_hash = MagicMock(return_value=result) + service = self._build_model_service_for_unit_test( + service_module="civitai_search_service", + service_class="CivitAISearchService", + download_available=True, + extract_sha256_from_metadata=lambda value: value.get("sha256", ""), + find_local_file_path=lambda filename, category=None: None, + get_filename_from_path=lambda value: str(value or "").replace( + "\\", "/" + ).rsplit("/", 1)[-1], + is_path_in_configured_model_roots=lambda path: True, + looks_like_model_file=lambda url, expected_filename="": str(url).endswith( + ".safetensors" + ), + normalize_category_to_model_type=lambda value: value, + normalize_sha256=lambda value: str(value or ""), + request_public_url=request_public_url, + resolve_civarchive_by_hash=resolve_civarchive_by_hash, + search_huggingface_for_file=MagicMock(return_value=None), + to_bool=to_bool, + ) + with tempfile.NamedTemporaryFile( + suffix=".safetensors", delete=False + ) as handle: + model_path = handle.name + request = AsyncMock() + request.json.return_value = { + "filename": "model.safetensors", + "resolved_path": model_path, + "sha256": expected_hash, + } + + try: + with patch( + "core.sources.civitai.get_model_info_by_hash", + return_value=None, + ): + api_response = await service.civitai_search(request) + finally: + os.unlink(model_path) + + self.assertEqual(api_response.status, 200) + resolve_civarchive_by_hash.assert_called_once() + request_public_url.assert_called_once() + self.assertEqual(request_public_url.call_args.args[1], mirror_url) + payload = json.loads(api_response.text) + self.assertEqual(payload["download_url"], mirror_url) + async def test_custom_url_service_rejects_unsafe_url(self): class TestUnsafeUrlError(Exception): pass diff --git a/tests/test_source_internals.py b/tests/test_source_internals.py index 9c5ef96..826c6ba 100644 --- a/tests/test_source_internals.py +++ b/tests/test_source_internals.py @@ -22,7 +22,7 @@ import requests # civarchive internals # --------------------------------------------------------------------------- from core.sources.civarchive import ( - _collect_download_urls_unified, + _collect_archive_download_urls, _fetch_remote_file_size_bytes, _normalize_archive_version, _normalize_download_url, @@ -560,9 +560,9 @@ class RemoteLinkIsDeadTests(unittest.TestCase): # =========================================================================== -# civarchive - _collect_download_urls_unified (new consolidated helper) +# civarchive - _collect_archive_download_urls (common collector adapter) # =========================================================================== -class CollectDownloadUrlsUnifiedTests(unittest.TestCase): +class CollectArchiveDownloadUrlsTests(unittest.TestCase): def _make_file_info(self, mirrors=None, download_url=None, download_urls=None, dead=False): @@ -581,14 +581,14 @@ class CollectDownloadUrlsUnifiedTests(unittest.TestCase): fi = self._make_file_info(mirrors=[ {"url": "https://civarchive.com/api/download/models/1.safetensors"}, ]) - result = _collect_download_urls_unified(fi) + result = _collect_archive_download_urls(fi) self.assertIn("https://civarchive.com/api/download/models/1.safetensors", result) def test_skips_dead_mirrors(self): fi = self._make_file_info(mirrors=[ {"url": "https://civarchive.com/api/download/models/1.safetensors", "isDead": True}, ]) - result = _collect_download_urls_unified(fi) + result = _collect_archive_download_urls(fi) self.assertEqual(result, []) def test_skip_file_if_dead_returns_empty(self): @@ -596,14 +596,14 @@ class CollectDownloadUrlsUnifiedTests(unittest.TestCase): download_url="https://civarchive.com/api/download/models/1.safetensors", dead=True, ) - result = _collect_download_urls_unified(fi, skip_file_if_dead=True) + result = _collect_archive_download_urls(fi, skip_file_if_dead=True) self.assertEqual(result, []) def test_check_download_urls_list_includes_extra_urls(self): fi = self._make_file_info(download_urls=[ "https://civarchive.com/api/download/models/2.safetensors" ]) - result = _collect_download_urls_unified(fi, check_download_urls_list=True, + result = _collect_archive_download_urls(fi, check_download_urls_list=True, skip_file_if_dead=True) self.assertIn("https://civarchive.com/api/download/models/2.safetensors", result) @@ -614,7 +614,7 @@ class CollectDownloadUrlsUnifiedTests(unittest.TestCase): {"url": civitai_url}, {"url": other_url}, ]) - result = _collect_download_urls_unified(fi, prioritize_civitai_last=True) + result = _collect_archive_download_urls(fi, prioritize_civitai_last=True) self.assertEqual(result[-1], civitai_url) self.assertEqual(result[0], other_url) @@ -624,23 +624,23 @@ class CollectDownloadUrlsUnifiedTests(unittest.TestCase): mirrors=[{"url": url}], download_url=url, ) - result = _collect_download_urls_unified(fi) + result = _collect_archive_download_urls(fi) self.assertEqual(result.count(url), 1) def test_empty_file_info_returns_empty(self): - self.assertEqual(_collect_download_urls_unified({}), []) + self.assertEqual(_collect_archive_download_urls({}), []) # =========================================================================== # civarchive - archive download URL collection strategy # =========================================================================== -class CollectDownloadUrlsStrategyTests(unittest.TestCase): +class CollectArchiveDownloadUrlsStrategyTests(unittest.TestCase): def test_civitai_url_sorted_last(self): civitai = "https://civitai.com/api/download/models/99.safetensors" archive = "https://civarchive.com/api/download/models/77.safetensors" fi = {"mirrors": [{"url": civitai}, {"url": archive}]} - result = _collect_download_urls_unified( + result = _collect_archive_download_urls( fi, prioritize_civitai_last=True, download_url_keys=("downloadUrl",), @@ -649,7 +649,7 @@ class CollectDownloadUrlsStrategyTests(unittest.TestCase): self.assertEqual(result[0], archive) def test_returns_list(self): - result = _collect_download_urls_unified( + result = _collect_archive_download_urls( {}, prioritize_civitai_last=True, download_url_keys=("downloadUrl",), @@ -660,7 +660,7 @@ class CollectDownloadUrlsStrategyTests(unittest.TestCase): # =========================================================================== # civarchive - normalized download URL collection strategy # =========================================================================== -class CollectNormalizedDownloadUrlsTests(unittest.TestCase): +class CollectNormalizedArchiveDownloadUrlsTests(unittest.TestCase): def test_skips_dead_file(self): fi = { @@ -668,7 +668,7 @@ class CollectNormalizedDownloadUrlsTests(unittest.TestCase): "downloadUrl": "https://civarchive.com/api/download/models/1.safetensors", } self.assertEqual( - _collect_download_urls_unified( + _collect_archive_download_urls( fi, skip_file_if_dead=True, check_download_urls_list=True, @@ -682,7 +682,7 @@ class CollectNormalizedDownloadUrlsTests(unittest.TestCase): "mirrors": [{"url": "https://civarchive.com/api/download/models/1.safetensors"}], "download_urls": ["https://civarchive.com/api/download/models/2.safetensors"], } - result = _collect_download_urls_unified( + result = _collect_archive_download_urls( fi, skip_file_if_dead=True, check_download_urls_list=True, @@ -696,7 +696,7 @@ class CollectNormalizedDownloadUrlsTests(unittest.TestCase): "mirrors": [{"url": dead_url, "isDead": True}], "download_urls": [dead_url], } - result = _collect_download_urls_unified( + result = _collect_archive_download_urls( fi, skip_file_if_dead=True, check_download_urls_list=True,