943 lines
31 KiB
Python
943 lines
31 KiB
Python
import asyncio
|
|
from types import SimpleNamespace
|
|
from unittest.mock import AsyncMock, MagicMock
|
|
|
|
import pytest
|
|
|
|
from core.contracts import ModelCatalogEntry, ModelMatch, SearchResult
|
|
from core.routes.context import RouteContext
|
|
from core.services.search_cache import SearchResultCache
|
|
from core.services.search_dependencies import SearchDependencies
|
|
from core.services.search_orchestrator import SearchOrchestrator
|
|
from core.services.search_providers import (
|
|
SearchCancelled,
|
|
SearchProviderRunner,
|
|
SearchRequest,
|
|
)
|
|
from core.type_utils import build_model_result
|
|
|
|
|
|
def _request(**overrides):
|
|
values = {
|
|
"filename": "model.safetensors",
|
|
"category": "checkpoints",
|
|
"base_model_context": "SDXL",
|
|
"progress_id": "search-1",
|
|
"progress_source": "civitai",
|
|
"force_search": False,
|
|
}
|
|
values.update(overrides)
|
|
request_data = values.get("data")
|
|
if isinstance(request_data, dict):
|
|
values.setdefault("model_id", request_data.get("model_id"))
|
|
values.setdefault("version_id", request_data.get("version_id"))
|
|
return SimpleNamespace(**values)
|
|
|
|
|
|
def _typed_search_request(**overrides):
|
|
values = {
|
|
"filename": "model.safetensors",
|
|
"category": "checkpoints",
|
|
"base_model_context": "SDXL",
|
|
"progress_id": "search-1",
|
|
"progress_source": "civitai",
|
|
"civitai_candidate_limit": 5,
|
|
"civarchive_candidate_limit": 10,
|
|
"is_urn": False,
|
|
"civitai_key": "",
|
|
"civitai_session_token": "",
|
|
"hf_token": "",
|
|
"brave_search_api_key": "",
|
|
"civitai_use_trpc_search": True,
|
|
"civitai_use_api_search": True,
|
|
"civitai_use_html_fallback": True,
|
|
"hf_use_api_search": True,
|
|
"hf_use_comfy_org_fallback": True,
|
|
"hf_use_brave_fallback": True,
|
|
"force_search": False,
|
|
"normalized_sources": frozenset({"civitai"}),
|
|
"sha256": "",
|
|
}
|
|
values.update(overrides)
|
|
return SearchRequest(**values)
|
|
|
|
|
|
def test_search_request_from_mapping_adapts_http_payload_once():
|
|
request = SearchRequest.from_mapping(
|
|
{
|
|
"filename": " model.safetensors ",
|
|
"sources": ["CivitAI"],
|
|
"civitai_candidate_limit": 100,
|
|
"is_urn": "true",
|
|
"model_id": " 12 ",
|
|
"version_id": 34,
|
|
"sha256": "A" * 64,
|
|
},
|
|
)
|
|
|
|
assert request.filename == "model.safetensors"
|
|
assert request.normalized_sources == frozenset({"civitai"})
|
|
assert request.civitai_candidate_limit == 20
|
|
assert request.is_urn is True
|
|
assert request.model_id == "12"
|
|
assert request.version_id == 34
|
|
assert request.sha256 == "a" * 64
|
|
|
|
|
|
def test_search_request_from_mapping_rejects_non_object_payloads():
|
|
with pytest.raises(TypeError, match="Search request must be an object"):
|
|
SearchRequest.from_mapping(
|
|
[],
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("payload", "message"),
|
|
[
|
|
({"filename": 123}, "SearchRequest filename must be a string"),
|
|
(
|
|
{"filename": "model.safetensors", "category": 123},
|
|
"SearchRequest category must be a string",
|
|
),
|
|
(
|
|
{"filename": "model.safetensors", "sources": {"civitai": True}},
|
|
"SearchRequest sources must be a string or array",
|
|
),
|
|
(
|
|
{"filename": "model.safetensors", "sources": ["civitai", 1]},
|
|
"SearchRequest sources must contain only strings",
|
|
),
|
|
(
|
|
{"filename": "model.safetensors", "sources": None},
|
|
"SearchRequest sources must be a string or array",
|
|
),
|
|
],
|
|
)
|
|
def test_search_request_from_mapping_rejects_malformed_http_fields(payload, message):
|
|
with pytest.raises(TypeError, match=message):
|
|
SearchRequest.from_mapping(
|
|
payload,
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("field_name", "value"),
|
|
[
|
|
("category", 123),
|
|
("civitai_key", 123),
|
|
("is_urn", "false"),
|
|
("force_search", {"enabled": True}),
|
|
("civitai_candidate_limit", "5"),
|
|
("normalized_sources", {"civitai"}),
|
|
("model_id", 1.5),
|
|
("version_id", False),
|
|
],
|
|
)
|
|
def test_search_request_rejects_invalid_normalized_field_types(field_name, value):
|
|
with pytest.raises((TypeError, ValueError), match=field_name):
|
|
_typed_search_request(**{field_name: value})
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"field_name",
|
|
["civitai_candidate_limit", "civarchive_candidate_limit"],
|
|
)
|
|
def test_search_request_from_mapping_rejects_invalid_limit_types(field_name):
|
|
with pytest.raises((TypeError, ValueError), match=field_name):
|
|
SearchRequest.from_mapping(
|
|
{"filename": "model.safetensors", field_name: "invalid"},
|
|
)
|
|
|
|
|
|
def _build_search_orchestrator():
|
|
tracker = MagicMock()
|
|
tracker.is_cancelled.return_value = False
|
|
extension = SimpleNamespace(
|
|
search_tracker=tracker,
|
|
search_result_timestamps={},
|
|
logger=MagicMock(),
|
|
)
|
|
|
|
def extract_sha256(metadata):
|
|
return metadata.get("sha256") or (
|
|
metadata.get("hashes") or {}
|
|
).get("SHA256")
|
|
|
|
context_values = {
|
|
"self": extension,
|
|
"asyncio": asyncio,
|
|
"CivArchiveSearchError": Exception,
|
|
"build_model_result": build_model_result,
|
|
"clear_civarchive_search_cache": MagicMock(),
|
|
"clear_civitai_search_cache": MagicMock(),
|
|
"clear_huggingface_search_cache": MagicMock(),
|
|
"clear_lora_manager_archive_search_cache": MagicMock(),
|
|
"extract_sha256_from_metadata": extract_sha256,
|
|
"format_size_bytes": lambda value, include_space=True: str(value),
|
|
"get_civitai_download_url": MagicMock(),
|
|
"get_popular_model_url": MagicMock(),
|
|
"reload_model_list": MagicMock(),
|
|
"reload_popular_databases": MagicMock(),
|
|
"resolve_civarchive_model_version": MagicMock(),
|
|
"resolve_urn": MagicMock(),
|
|
"search_civarchive_for_file": MagicMock(),
|
|
"search_civitai": MagicMock(),
|
|
"search_civitai_for_file": MagicMock(),
|
|
"search_huggingface_for_file": MagicMock(),
|
|
"search_local_matches_by_hash": MagicMock(return_value=[]),
|
|
"search_lora_manager_archive_for_file": MagicMock(),
|
|
"search_model_list": MagicMock(),
|
|
"web": SimpleNamespace(
|
|
json_response=lambda payload, status=200: SimpleNamespace(
|
|
payload=payload,
|
|
status=status,
|
|
)
|
|
),
|
|
}
|
|
return SearchOrchestrator(RouteContext(context_values)), context_values
|
|
|
|
|
|
class _StaticSearchRunner:
|
|
def __init__(self, source_results, source_found):
|
|
self.source_results = source_results
|
|
self.source_found = source_found
|
|
|
|
def create_search_tasks(self, request):
|
|
async def search_task():
|
|
return self.source_results, self.source_found
|
|
|
|
return [search_task()]
|
|
|
|
def raise_if_search_cancelled(self, request, source=""):
|
|
return None
|
|
|
|
|
|
def test_search_result_cache_reuses_timestamp_for_same_result():
|
|
timestamps = {}
|
|
cache = SearchResultCache(timestamps)
|
|
|
|
first = cache.stamp_results(
|
|
{"civitai": {"name": "Model", "download_url": "https://example.test/model"}},
|
|
force_search=False,
|
|
)
|
|
second = cache.stamp_results(
|
|
{"civitai": {"name": "Model", "download_url": "https://example.test/model"}},
|
|
force_search=False,
|
|
)
|
|
|
|
assert first["civitai"]["searched_at"] == second["civitai"]["searched_at"]
|
|
assert len(timestamps) == 1
|
|
|
|
|
|
def test_search_result_cache_stamps_nested_results_and_preserves_empty_values():
|
|
cache = SearchResultCache({})
|
|
payload = {
|
|
"popular": [
|
|
{"filename": "one.safetensors", "url": "https://example.test/one"},
|
|
{"filename": "two.safetensors", "url": "https://example.test/two"},
|
|
],
|
|
"civitai": None,
|
|
"source_errors": {"civitai": "unavailable"},
|
|
}
|
|
|
|
stamped = cache.stamp_results(payload, force_search=False)
|
|
|
|
assert all(item["searched_at"] for item in stamped["popular"])
|
|
assert stamped["civitai"] is None
|
|
assert stamped["source_errors"] == {"civitai": "unavailable"}
|
|
|
|
|
|
def test_search_provider_runner_reports_progress_and_result():
|
|
owner = SimpleNamespace(
|
|
logger=MagicMock(),
|
|
search_tracker=MagicMock(),
|
|
format_log_fields=MagicMock(return_value="file=model.safetensors"),
|
|
)
|
|
owner.search_tracker.is_cancelled.return_value = False
|
|
runner = SearchProviderRunner(owner)
|
|
|
|
result = runner.run_source_search(
|
|
_request(),
|
|
"civitai",
|
|
lambda: ({"civitai": {"name": "Model"}}, True),
|
|
initial_stage="query",
|
|
initial_message="Querying CivitAI",
|
|
initial_percent=30,
|
|
)
|
|
|
|
assert result == ({"civitai": {"name": "Model"}}, True)
|
|
progress_calls = owner.search_tracker.update.call_args_list
|
|
assert progress_calls[0].args[:5] == (
|
|
"search-1",
|
|
"civitai",
|
|
"query",
|
|
"Querying CivitAI",
|
|
30,
|
|
)
|
|
assert progress_calls[-1].args[:5] == (
|
|
"search-1",
|
|
"civitai",
|
|
"done",
|
|
"CivitAI checked",
|
|
92,
|
|
)
|
|
|
|
|
|
def test_search_provider_runner_raises_cancellation_before_provider_call():
|
|
owner = SimpleNamespace(search_tracker=MagicMock())
|
|
owner.search_tracker.is_cancelled.return_value = True
|
|
runner = SearchProviderRunner(owner)
|
|
|
|
with pytest.raises(SearchCancelled):
|
|
runner.raise_if_search_cancelled(_request(), "civitai")
|
|
|
|
|
|
def test_search_provider_runner_retries_without_base_model_context():
|
|
owner = SimpleNamespace(
|
|
logger=MagicMock(),
|
|
search_tracker=MagicMock(),
|
|
format_log_fields=MagicMock(return_value="file=model.safetensors"),
|
|
log_search_result=MagicMock(),
|
|
)
|
|
owner.search_tracker.is_cancelled.return_value = False
|
|
runner = SearchProviderRunner(owner)
|
|
calls = []
|
|
|
|
def search_fn(base_model_context, progress_callback):
|
|
calls.append((base_model_context, progress_callback))
|
|
if len(calls) == 1:
|
|
return None
|
|
return SearchResult(source="civitai", name="Fallback model")
|
|
|
|
result = runner.execute_search_with_fallback(
|
|
_request(),
|
|
"civitai",
|
|
search_fn,
|
|
"CivitAI",
|
|
)
|
|
|
|
assert [call[0] for call in calls] == ["SDXL", None]
|
|
assert result.extra_value("any_model_match") is True
|
|
assert result.extra_value("base_model_fallback") is True
|
|
assert result.extra_value("requested_base_model") == "SDXL"
|
|
|
|
|
|
def test_search_provider_runner_searches_local_sources():
|
|
owner = SimpleNamespace(
|
|
logger=MagicMock(),
|
|
search_tracker=MagicMock(),
|
|
format_log_fields=MagicMock(return_value="file=model.safetensors"),
|
|
build_model_result=build_model_result,
|
|
get_popular_model_url=MagicMock(
|
|
return_value=ModelCatalogEntry(
|
|
filename="model.safetensors",
|
|
download_url="https://example.test/model",
|
|
)
|
|
),
|
|
search_model_list=MagicMock(
|
|
return_value=SearchResult(
|
|
source="model_list",
|
|
filename="model.safetensors",
|
|
size=123,
|
|
confidence=95,
|
|
)
|
|
),
|
|
log_search_result=MagicMock(),
|
|
)
|
|
owner.search_tracker.is_cancelled.return_value = False
|
|
runner = SearchProviderRunner(owner)
|
|
|
|
result, found = runner.search_local_sources(_request(is_urn=False))
|
|
|
|
assert found is True
|
|
assert result["popular"].size == 123
|
|
assert result["popular"].download_url == "https://example.test/model"
|
|
assert result["model_list"].confidence == 95
|
|
|
|
|
|
def test_search_provider_runner_scales_provider_progress():
|
|
owner = SimpleNamespace(search_tracker=MagicMock())
|
|
owner.search_tracker.is_cancelled.return_value = False
|
|
runner = SearchProviderRunner(owner)
|
|
|
|
callback = runner.make_source_progress_callback(
|
|
_request(),
|
|
"civitai",
|
|
20,
|
|
80,
|
|
)
|
|
callback({"stage": "query", "message": "Searching", "percent": 50, "source": "ignored", "items": 2})
|
|
|
|
update_call = owner.search_tracker.update.call_args
|
|
assert update_call.args[:5] == (
|
|
"search-1",
|
|
"civitai",
|
|
"query",
|
|
"Searching",
|
|
50.0,
|
|
)
|
|
assert update_call.kwargs == {"status": "running", "items": 2}
|
|
|
|
|
|
def test_search_provider_runner_handles_civarchive_error():
|
|
class ProviderError(Exception):
|
|
pass
|
|
|
|
owner = SimpleNamespace(
|
|
CivArchiveSearchError=ProviderError,
|
|
logger=MagicMock(),
|
|
search_tracker=MagicMock(),
|
|
format_log_fields=MagicMock(return_value="file=model.safetensors"),
|
|
)
|
|
owner.search_tracker.is_cancelled.return_value = False
|
|
runner = SearchProviderRunner(owner)
|
|
|
|
result = runner.run_source_search(
|
|
_request(),
|
|
"civarchive",
|
|
lambda: (_ for _ in ()).throw(ProviderError("provider unavailable")),
|
|
error_handlers={
|
|
ProviderError: lambda error: (
|
|
{"civarchive": None, "source_errors": {"civarchive": str(error)}},
|
|
False,
|
|
)
|
|
},
|
|
)
|
|
|
|
assert result == (
|
|
{"civarchive": None, "source_errors": {"civarchive": "provider unavailable"}},
|
|
False,
|
|
)
|
|
|
|
|
|
def test_search_provider_runner_exposes_civarchive_status_metadata():
|
|
from core.sources.civarchive import CivArchiveSearchError
|
|
|
|
owner = SimpleNamespace(
|
|
CivArchiveSearchError=CivArchiveSearchError,
|
|
logger=MagicMock(),
|
|
search_tracker=MagicMock(),
|
|
format_log_fields=MagicMock(return_value="file=model.safetensors"),
|
|
log_search_result=MagicMock(),
|
|
search_civarchive_for_file=MagicMock(
|
|
side_effect=CivArchiveSearchError(
|
|
"HTTP 522",
|
|
code="provider_unavailable",
|
|
http_status=522,
|
|
retryable=True,
|
|
)
|
|
),
|
|
)
|
|
owner.search_tracker.is_cancelled.return_value = False
|
|
runner = SearchProviderRunner(owner)
|
|
|
|
result = runner.search_civarchive_source_task(
|
|
_request(is_urn=False, data={}, civarchive_candidate_limit=5)
|
|
)
|
|
|
|
assert result[0]["source_status"]["civarchive"] == {
|
|
"state": "unavailable",
|
|
"code": "provider_unavailable",
|
|
"retryable": True,
|
|
"http_status": 522,
|
|
"message": (
|
|
"CivArchive may be overloaded or temporarily unavailable. "
|
|
"Please try again."
|
|
),
|
|
}
|
|
warning_message = owner.logger.warning.call_args.args[0]
|
|
assert warning_message == (
|
|
"CivArchive search failed: code=provider_unavailable "
|
|
"http_status=522 retryable=yes error_type=CivArchiveSearchError "
|
|
"message=CivArchive may be overloaded or temporarily unavailable. "
|
|
"Please try again."
|
|
)
|
|
|
|
|
|
def test_search_provider_runner_resolves_civitai_urn():
|
|
owner = SimpleNamespace(
|
|
CivArchiveSearchError=Exception,
|
|
logger=MagicMock(),
|
|
search_tracker=MagicMock(),
|
|
format_log_fields=MagicMock(return_value="ids=10@20"),
|
|
resolve_urn=MagicMock(
|
|
return_value={
|
|
"model_name": "URN Model",
|
|
"version_name": "v1",
|
|
"expected_filename": "urn.safetensors",
|
|
"base_model": "SDXL",
|
|
"files": [{"name": "urn.safetensors", "sha256": "a" * 64}],
|
|
}
|
|
),
|
|
get_civitai_download_url=MagicMock(return_value="https://example.test/download"),
|
|
build_model_result=MagicMock(side_effect=build_model_result),
|
|
)
|
|
owner.search_tracker.is_cancelled.return_value = False
|
|
owner.log_search_result = MagicMock()
|
|
runner = SearchProviderRunner(owner)
|
|
|
|
result, found = runner.search_civitai_source_task(
|
|
_request(
|
|
data={"model_id": 10, "version_id": 20},
|
|
is_urn=True,
|
|
base_model_context="",
|
|
)
|
|
)
|
|
|
|
assert found is True
|
|
assert result["civitai"].model_id == 10
|
|
assert result["civitai"].version_id == 20
|
|
assert result["civitai"].name == "URN Model"
|
|
assert result["civitai"].version_name == "v1"
|
|
assert result["civitai"].filename == "urn.safetensors"
|
|
assert result["civitai"].model_type == "checkpoints"
|
|
assert result["civitai"].download_url == "https://example.test/download"
|
|
assert result["civitai"].url == "https://civitai.com/models/10?modelVersionId=20"
|
|
assert result["civitai"].size is None
|
|
assert result["civitai"].base_model == "SDXL"
|
|
assert result["civitai"].match_type == "exact"
|
|
assert result["civitai"].confidence == 100.0
|
|
assert result["civitai"].sha256 == "a" * 64
|
|
assert result["civitai"].hashes == {}
|
|
|
|
|
|
def test_search_provider_runner_keeps_first_file_when_expected_name_is_missing():
|
|
owner = SimpleNamespace(
|
|
CivArchiveSearchError=Exception,
|
|
logger=MagicMock(),
|
|
search_tracker=MagicMock(),
|
|
format_log_fields=MagicMock(return_value="ids=10@20"),
|
|
resolve_urn=MagicMock(
|
|
return_value={
|
|
"model_name": "URN Model",
|
|
"version_name": "v1",
|
|
"expected_filename": "missing.safetensors",
|
|
"base_model": "SDXL",
|
|
"files": [
|
|
{"name": "fallback.safetensors", "sha256": "a" * 64},
|
|
{
|
|
"name": "primary.safetensors",
|
|
"primary": True,
|
|
"type": "Model",
|
|
"sha256": "b" * 64,
|
|
},
|
|
],
|
|
}
|
|
),
|
|
get_civitai_download_url=MagicMock(return_value="https://example.test/download"),
|
|
build_model_result=MagicMock(side_effect=build_model_result),
|
|
)
|
|
owner.search_tracker.is_cancelled.return_value = False
|
|
owner.log_search_result = MagicMock()
|
|
runner = SearchProviderRunner(owner)
|
|
|
|
result, found = runner.search_civitai_source_task(
|
|
_request(
|
|
data={"model_id": 10, "version_id": 20},
|
|
is_urn=True,
|
|
base_model_context="",
|
|
)
|
|
)
|
|
|
|
assert found is True
|
|
assert result["civitai"].sha256 == "a" * 64
|
|
|
|
|
|
def test_search_provider_runner_falls_back_for_civitai_urn_without_ids():
|
|
owner = SimpleNamespace(
|
|
logger=MagicMock(),
|
|
search_tracker=MagicMock(),
|
|
format_log_fields=MagicMock(return_value="file=model.safetensors"),
|
|
search_civitai=MagicMock(
|
|
return_value=[
|
|
SearchResult(
|
|
source="civitai",
|
|
model_id=30,
|
|
version_id=40,
|
|
name="Fallback model",
|
|
filename="model.safetensors",
|
|
model_type="Checkpoint",
|
|
download_url="https://example.test/fallback",
|
|
url="https://example.test/model/30",
|
|
size=456,
|
|
base_model="SDXL",
|
|
tags=["fallback"],
|
|
)
|
|
]
|
|
),
|
|
build_model_result=MagicMock(
|
|
side_effect=lambda source, **fields: {"source": source, **fields}
|
|
),
|
|
log_search_result=MagicMock(),
|
|
)
|
|
owner.search_tracker.is_cancelled.return_value = False
|
|
runner = SearchProviderRunner(owner)
|
|
|
|
result, found = runner.search_civitai_source_task(
|
|
_request(data={"model_id": None, "version_id": None}, is_urn=True)
|
|
)
|
|
|
|
assert found is True
|
|
assert isinstance(result["civitai"], SearchResult)
|
|
assert result["civitai"].model_id == 30
|
|
assert result["civitai"].download_url == "https://example.test/fallback"
|
|
owner.search_civitai.assert_called_once_with(
|
|
"model.safetensors",
|
|
model_type="checkpoints",
|
|
)
|
|
|
|
|
|
def test_search_provider_runner_resolves_civarchive_urn():
|
|
owner = SimpleNamespace(
|
|
CivArchiveSearchError=Exception,
|
|
logger=MagicMock(),
|
|
search_tracker=MagicMock(),
|
|
format_log_fields=MagicMock(return_value="ids=10@20"),
|
|
resolve_civarchive_model_version=MagicMock(
|
|
return_value=SearchResult(
|
|
source="civarchive",
|
|
name="Archived model",
|
|
version_id=20,
|
|
)
|
|
),
|
|
log_search_result=MagicMock(),
|
|
)
|
|
owner.search_tracker.is_cancelled.return_value = False
|
|
runner = SearchProviderRunner(owner)
|
|
|
|
result, found = runner.search_civarchive_source_task(
|
|
_request(
|
|
data={"model_id": 10, "version_id": 20},
|
|
is_urn=True,
|
|
)
|
|
)
|
|
|
|
assert found is True
|
|
assert result["civarchive"].name == "Archived model"
|
|
assert result["civarchive"].version_id == 20
|
|
owner.resolve_civarchive_model_version.assert_called_once_with(
|
|
10,
|
|
20,
|
|
query="model.safetensors",
|
|
)
|
|
|
|
|
|
def test_search_provider_runner_searches_lora_manager_archive():
|
|
owner = SimpleNamespace(
|
|
logger=MagicMock(),
|
|
search_tracker=MagicMock(),
|
|
format_log_fields=MagicMock(return_value="file=model.safetensors"),
|
|
search_lora_manager_archive_for_file=MagicMock(
|
|
return_value=SearchResult(source="lora_manager_archive", name="Archived LoRA")
|
|
),
|
|
log_search_result=MagicMock(),
|
|
)
|
|
owner.search_tracker.is_cancelled.return_value = False
|
|
runner = SearchProviderRunner(owner)
|
|
|
|
result, found = runner.search_lora_manager_archive_source_task(
|
|
_request(base_model_context="")
|
|
)
|
|
|
|
assert found is True
|
|
assert result["lora_manager_archive"].name == "Archived LoRA"
|
|
owner.search_lora_manager_archive_for_file.assert_called_once()
|
|
|
|
|
|
def test_search_provider_runner_skips_name_fallback_for_sha256():
|
|
owner = SimpleNamespace(
|
|
logger=MagicMock(),
|
|
search_tracker=MagicMock(),
|
|
format_log_fields=MagicMock(return_value="file=model.safetensors"),
|
|
log_search_result=MagicMock(),
|
|
)
|
|
owner.search_tracker.is_cancelled.return_value = False
|
|
runner = SearchProviderRunner(owner)
|
|
search_fn = MagicMock(return_value=None)
|
|
|
|
result = runner.execute_search_with_fallback(
|
|
_request(sha256="a" * 64),
|
|
"lora_manager_archive",
|
|
search_fn,
|
|
"LoRA archive",
|
|
)
|
|
|
|
assert result is None
|
|
search_fn.assert_called_once()
|
|
|
|
|
|
def test_search_dependencies_require_extension_context():
|
|
with pytest.raises(KeyError, match="Missing route dependency: self"):
|
|
SearchDependencies.from_context(RouteContext({}))
|
|
|
|
|
|
def test_search_dependencies_require_core_provider_dependencies():
|
|
extension = SimpleNamespace(
|
|
logger=MagicMock(),
|
|
search_result_timestamps={},
|
|
search_tracker=MagicMock(),
|
|
)
|
|
|
|
with pytest.raises(KeyError, match="Missing route dependency: asyncio"):
|
|
SearchDependencies.from_context(RouteContext({"self": extension}))
|
|
|
|
|
|
def test_search_orchestrator_formats_log_values_and_result_details():
|
|
orchestrator, _ = _build_search_orchestrator()
|
|
|
|
assert orchestrator.format_log_value(None) is None
|
|
assert orchestrator.format_log_value("") is None
|
|
assert orchestrator.format_log_value(False) == "no"
|
|
assert orchestrator.format_log_value(["civitai", "huggingface"]) == (
|
|
"civitai,huggingface"
|
|
)
|
|
assert orchestrator.format_log_value('model "one"') == '"model \\"one\\""'
|
|
assert orchestrator.format_log_value("plain") == "plain"
|
|
assert orchestrator.format_log_fields(
|
|
enabled=False,
|
|
tags=["one", "two"],
|
|
missing=None,
|
|
) == "enabled=no tags=one,two"
|
|
assert orchestrator.format_result_details([{"name": "one"}]) == "count=1"
|
|
assert orchestrator.format_result_details(
|
|
None,
|
|
{"model_id": 10, "files_count": 2},
|
|
) == "result=none ids=10 files=2"
|
|
assert "ids=10@20" in orchestrator.format_result_details(
|
|
{
|
|
"name": "Model",
|
|
"filename": "model.safetensors",
|
|
"model_id": 10,
|
|
"version_id": 20,
|
|
"size": 1024,
|
|
}
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_search_orchestrator_force_search_refreshes_selected_caches():
|
|
orchestrator, dependencies = _build_search_orchestrator()
|
|
orchestrator.provider_runner = _StaticSearchRunner({}, False)
|
|
request = SimpleNamespace(
|
|
json=AsyncMock(
|
|
return_value={
|
|
"filename": "model.safetensors",
|
|
"sources": [
|
|
"local",
|
|
"huggingface",
|
|
"civitai",
|
|
"civarchive",
|
|
"lora_manager_archive",
|
|
],
|
|
"force_search": True,
|
|
"progress_id": "force-search",
|
|
}
|
|
)
|
|
)
|
|
|
|
response = await orchestrator.search_sources(request)
|
|
|
|
assert response.status == 200
|
|
assert response.payload["found"] is False
|
|
dependencies["reload_popular_databases"].assert_called_once_with()
|
|
dependencies["reload_model_list"].assert_called_once_with()
|
|
dependencies["clear_huggingface_search_cache"].assert_called_once_with()
|
|
dependencies["clear_civitai_search_cache"].assert_called_once_with()
|
|
dependencies["clear_civarchive_search_cache"].assert_called_once_with()
|
|
dependencies[
|
|
"clear_lora_manager_archive_search_cache"
|
|
].assert_called_once_with()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_search_orchestrator_preserves_provider_status_metadata():
|
|
orchestrator, _dependencies = _build_search_orchestrator()
|
|
orchestrator.provider_runner = _StaticSearchRunner(
|
|
{
|
|
"civarchive": None,
|
|
"source_errors": {
|
|
"civarchive": "CivArchive search failed: HTTP 522"
|
|
},
|
|
"source_status": {
|
|
"civarchive": {
|
|
"state": "unavailable",
|
|
"code": "provider_unavailable",
|
|
"retryable": True,
|
|
"http_status": 522,
|
|
"message": (
|
|
"CivArchive may be overloaded or temporarily unavailable. "
|
|
"Please try again."
|
|
),
|
|
}
|
|
},
|
|
},
|
|
False,
|
|
)
|
|
request = SimpleNamespace(
|
|
json=AsyncMock(
|
|
return_value={
|
|
"filename": "model.safetensors",
|
|
"sources": ["civarchive"],
|
|
"progress_id": "provider-status-search",
|
|
}
|
|
)
|
|
)
|
|
|
|
response = await orchestrator.search_sources(request)
|
|
|
|
assert response.status == 200
|
|
assert response.payload["found"] is False
|
|
assert response.payload["source_errors"] == {
|
|
"civarchive": "CivArchive search failed: HTTP 522"
|
|
}
|
|
assert response.payload["source_status"]["civarchive"]["code"] == (
|
|
"provider_unavailable"
|
|
)
|
|
assert response.payload["source_status"]["civarchive"]["retryable"] is True
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_search_orchestrator_returns_cancelled_response():
|
|
orchestrator, dependencies = _build_search_orchestrator()
|
|
dependencies["self"].search_tracker.is_cancelled.return_value = True
|
|
request = SimpleNamespace(
|
|
json=AsyncMock(
|
|
return_value={
|
|
"filename": "model.safetensors",
|
|
"sources": ["civitai"],
|
|
"progress_id": "cancelled-search",
|
|
}
|
|
)
|
|
)
|
|
|
|
response = await orchestrator.search_sources(request)
|
|
|
|
assert response.status == 200
|
|
assert response.payload["cancelled"] is True
|
|
assert response.payload["found"] is False
|
|
assert response.payload["searched_sources"] == ["civitai"]
|
|
cancelled_updates = [
|
|
call
|
|
for call in dependencies["self"].search_tracker.update.call_args_list
|
|
if len(call.args) >= 3 and call.args[2] == "cancelled"
|
|
]
|
|
assert cancelled_updates
|
|
assert cancelled_updates[-1].kwargs == {
|
|
"status": "cancelled",
|
|
"cancelled": True,
|
|
}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_search_orchestrator_rejects_invalid_request_payload():
|
|
orchestrator, dependencies = _build_search_orchestrator()
|
|
request = SimpleNamespace(
|
|
json=AsyncMock(side_effect=ValueError("invalid request"))
|
|
)
|
|
|
|
response = await orchestrator.search_sources(request)
|
|
|
|
assert response.status == 400
|
|
assert response.payload == {
|
|
"error": "Filename is required for non-URN, or model_id+version_id for URN"
|
|
}
|
|
dependencies["self"].search_tracker.update.assert_not_called()
|
|
dependencies["self"].logger.exception.assert_not_called()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
"payload",
|
|
[
|
|
{"filename": "model.safetensors", "category": 123},
|
|
{"filename": "model.safetensors", "sources": {"civitai": True}},
|
|
],
|
|
)
|
|
async def test_search_orchestrator_returns_bad_request_for_contract_errors(payload):
|
|
orchestrator, dependencies = _build_search_orchestrator()
|
|
request = SimpleNamespace(json=AsyncMock(return_value=payload))
|
|
|
|
response = await orchestrator.search_sources(request)
|
|
|
|
assert response.status == 400
|
|
assert response.payload["error"]
|
|
dependencies["self"].search_tracker.update.assert_not_called()
|
|
dependencies["self"].logger.exception.assert_not_called()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_search_orchestrator_deduplicates_local_hash_matches_and_keeps_errors():
|
|
orchestrator, dependencies = _build_search_orchestrator()
|
|
first_hash = "a" * 64
|
|
failing_hash = "b" * 64
|
|
|
|
def lookup_by_hash(sha256, **kwargs):
|
|
if sha256 == first_hash:
|
|
return [
|
|
ModelMatch.from_mapping(
|
|
{
|
|
"model": {"path": r"C:\\Models\\Local.safetensors"},
|
|
"filename": "Local.safetensors",
|
|
}
|
|
),
|
|
ModelMatch.from_mapping(
|
|
{
|
|
"model": {"path": r"c:\\models\\local.safetensors"},
|
|
"filename": "local.safetensors",
|
|
}
|
|
),
|
|
]
|
|
raise RuntimeError("local index unavailable")
|
|
|
|
dependencies["search_local_matches_by_hash"].side_effect = lookup_by_hash
|
|
orchestrator.provider_runner = _StaticSearchRunner(
|
|
{
|
|
"civitai": [
|
|
SearchResult(
|
|
source="civitai",
|
|
filename="remote.safetensors",
|
|
hashes={"SHA256": first_hash},
|
|
),
|
|
SearchResult(
|
|
source="civitai",
|
|
filename="without-hash.safetensors",
|
|
),
|
|
SearchResult(
|
|
source="civitai",
|
|
filename="failing.safetensors",
|
|
sha256=failing_hash,
|
|
),
|
|
],
|
|
"source_errors": {"civarchive": "provider unavailable"},
|
|
},
|
|
True,
|
|
)
|
|
request = SimpleNamespace(
|
|
json=AsyncMock(
|
|
return_value={
|
|
"filename": "remote.safetensors",
|
|
"category": "checkpoints",
|
|
"sources": ["civitai"],
|
|
"progress_id": "hash-search",
|
|
}
|
|
)
|
|
)
|
|
|
|
response = await orchestrator.search_sources(request)
|
|
|
|
assert response.status == 200
|
|
assert response.payload["found"] is True
|
|
assert response.payload["source_errors"] == {
|
|
"civarchive": "provider unavailable"
|
|
}
|
|
assert response.payload["source_status"] == {}
|
|
assert len(response.payload["local_hash_matches"]) == 1
|
|
assert response.payload["local_hash_matches"][0]["hash_lookup_sha256"] == (
|
|
first_hash
|
|
)
|
|
assert response.payload["local_hash_matches"][0]["hash_lookup_source"] == (
|
|
"civitai"
|
|
)
|
|
assert dependencies["search_local_matches_by_hash"].call_count == 2
|
|
dependencies["self"].logger.warning.assert_called_once()
|