438 lines
13 KiB
Python
438 lines
13 KiB
Python
import asyncio
|
|
import json
|
|
import os
|
|
from types import SimpleNamespace
|
|
from unittest.mock import AsyncMock, MagicMock
|
|
|
|
import pytest
|
|
from aiohttp import web
|
|
|
|
from core.contracts import (
|
|
MissingModel,
|
|
ModelCatalogEntry,
|
|
ModelMatch,
|
|
Resolution,
|
|
SearchResult,
|
|
WorkflowAnalysisResult,
|
|
)
|
|
from core.path_utils import get_filename_from_path
|
|
from core.routes.context import RouteContext
|
|
from core.routes.workflow_analysis import register_workflow_analysis_routes
|
|
from core.type_utils import to_bool
|
|
|
|
|
|
class _Routes:
|
|
def __init__(self):
|
|
self.handlers = {}
|
|
|
|
def get(self, path):
|
|
def decorator(handler):
|
|
self.handlers[("GET", path)] = handler
|
|
return handler
|
|
|
|
return decorator
|
|
|
|
def post(self, path):
|
|
def decorator(handler):
|
|
self.handlers[("POST", path)] = handler
|
|
return handler
|
|
|
|
return decorator
|
|
|
|
|
|
def _json_api_endpoint(_name, **_kwargs):
|
|
return lambda handler: handler
|
|
|
|
|
|
def _request(payload):
|
|
return SimpleNamespace(json=AsyncMock(return_value=payload))
|
|
|
|
|
|
def _analysis_result(missing_models=(), **counts):
|
|
typed_missing = tuple(
|
|
item
|
|
if isinstance(item, MissingModel)
|
|
else MissingModel.from_mapping(item)
|
|
for item in missing_models
|
|
)
|
|
return WorkflowAnalysisResult(
|
|
missing_models=typed_missing,
|
|
total_resolved=counts.get("total_resolved", 0),
|
|
total_missing=counts.get("total_missing", len(typed_missing)),
|
|
total_models_analyzed=counts.get("total_models_analyzed", len(typed_missing)),
|
|
)
|
|
|
|
|
|
def _build_routes(overrides=None):
|
|
routes = _Routes()
|
|
extension = SimpleNamespace(
|
|
analysis_progress=MagicMock(),
|
|
_update_analysis_progress=MagicMock(),
|
|
logger=MagicMock(),
|
|
)
|
|
values = {
|
|
"analyze_and_find_matches": MagicMock(
|
|
return_value=WorkflowAnalysisResult()
|
|
),
|
|
"apply_resolution": MagicMock(side_effect=lambda workflow, _resolutions: workflow),
|
|
"asyncio": asyncio,
|
|
"download_available": False,
|
|
"fetch_remote_file_size_cached": MagicMock(return_value=123),
|
|
"get_filename_from_path": get_filename_from_path,
|
|
"get_popular_model_url": MagicMock(return_value=None),
|
|
"get_progress_response": MagicMock(return_value={"status": "running"}),
|
|
"invalidate_local_hash_match_cache": MagicMock(),
|
|
"json_api_endpoint": _json_api_endpoint,
|
|
"routes": routes,
|
|
"search_local_matches": MagicMock(return_value=[]),
|
|
"search_model_list": MagicMock(return_value=None),
|
|
"should_skip_existing_custom_node_reference": MagicMock(
|
|
return_value=False
|
|
),
|
|
"self": extension,
|
|
"to_bool": to_bool,
|
|
"web": web,
|
|
}
|
|
if overrides:
|
|
values.update(overrides)
|
|
register_workflow_analysis_routes(RouteContext(values))
|
|
return routes.handlers, values
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
("payload", "message"),
|
|
[
|
|
({}, "Workflow JSON is required"),
|
|
({"workflow": []}, "Workflow JSON must be an object"),
|
|
],
|
|
)
|
|
async def test_analyze_route_validates_workflow_payload(payload, message):
|
|
handlers, _ = _build_routes()
|
|
response = await handlers[("POST", "/model_resolver/analyze")](
|
|
_request(payload)
|
|
)
|
|
|
|
assert response.status == 400
|
|
assert json.loads(response.text) == {"error": message}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_analyze_route_rejects_non_text_analysis_id():
|
|
handlers, values = _build_routes()
|
|
response = await handlers[("POST", "/model_resolver/analyze")](
|
|
_request({"workflow": {"nodes": []}, "analysis_id": 123})
|
|
)
|
|
|
|
assert response.status == 400
|
|
assert "Workflow analysis request analysis_id must be a string" in (
|
|
json.loads(response.text)["error"]
|
|
)
|
|
values["analyze_and_find_matches"].assert_not_called()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_analyze_route_filters_references_and_uses_local_download_sources():
|
|
missing_models = [
|
|
{
|
|
"original_path": "hf/model.safetensors",
|
|
"matches": [],
|
|
"workflow_url": "https://huggingface.co/example/model.safetensors",
|
|
"workflow_model_url": "https://huggingface.co/example/model",
|
|
"workflow_directory": "checkpoints",
|
|
},
|
|
{
|
|
"original_path": "popular.safetensors",
|
|
"matches": [],
|
|
},
|
|
{
|
|
"original_path": "listed.safetensors",
|
|
"matches": [],
|
|
},
|
|
{
|
|
"original_path": "already-local.safetensors",
|
|
"matches": [
|
|
{
|
|
"model": {
|
|
"path": "models/already-local.safetensors",
|
|
"filename": "already-local.safetensors",
|
|
"category": "checkpoints",
|
|
},
|
|
"confidence": 100,
|
|
}
|
|
],
|
|
},
|
|
{
|
|
"original_path": "installed-custom-node.safetensors",
|
|
"matches": [],
|
|
},
|
|
]
|
|
|
|
def analyze(_workflow, _threshold, _limit, progress_callback, **_kwargs):
|
|
progress_callback({"stage": "matching", "current": 1, "total": 5})
|
|
return _analysis_result(missing_models)
|
|
|
|
def popular_url(filename):
|
|
if filename == "popular.safetensors":
|
|
return ModelCatalogEntry(
|
|
filename=filename,
|
|
url="https://example.com/popular.safetensors",
|
|
model_type="checkpoint",
|
|
directory="checkpoints",
|
|
size=100,
|
|
)
|
|
return None
|
|
|
|
def model_list(filename, exact_only=False):
|
|
assert exact_only is True
|
|
if filename == "popular.safetensors":
|
|
return SearchResult(
|
|
source="model_list",
|
|
filename=filename,
|
|
size=200,
|
|
)
|
|
if filename == "listed.safetensors":
|
|
return SearchResult(
|
|
source="model_list",
|
|
url="https://example.com/listed.safetensors",
|
|
filename="listed.safetensors",
|
|
name="Listed model",
|
|
model_type="checkpoint",
|
|
size=300,
|
|
match_type="exact",
|
|
confidence=100,
|
|
extra={"directory": "checkpoints"},
|
|
)
|
|
return None
|
|
|
|
handlers, values = _build_routes(
|
|
{
|
|
"analyze_and_find_matches": MagicMock(side_effect=analyze),
|
|
"download_available": True,
|
|
"get_popular_model_url": MagicMock(side_effect=popular_url),
|
|
"search_model_list": MagicMock(side_effect=model_list),
|
|
"should_skip_existing_custom_node_reference": MagicMock(
|
|
side_effect=lambda item: item.original_path.startswith(
|
|
"installed-custom"
|
|
)
|
|
),
|
|
}
|
|
)
|
|
response = await handlers[("POST", "/model_resolver/analyze")](
|
|
_request(
|
|
{
|
|
"workflow": {"nodes": []},
|
|
"analysis_id": "analysis-1",
|
|
"force_rescan": "true",
|
|
}
|
|
)
|
|
)
|
|
|
|
body = json.loads(response.text)
|
|
assert response.status == 200
|
|
assert body["total_missing"] == 4
|
|
assert all(
|
|
item["original_path"] != "installed-custom-node.safetensors"
|
|
for item in body["missing_models"]
|
|
)
|
|
sources = {
|
|
item["original_path"]: item["download_source"]
|
|
for item in body["missing_models"]
|
|
if "download_source" in item
|
|
}
|
|
assert sources["hf/model.safetensors"]["source"] == "huggingface"
|
|
assert sources["hf/model.safetensors"]["size"] == 123
|
|
assert sources["popular.safetensors"]["source"] == "popular"
|
|
assert sources["popular.safetensors"]["size"] == 200
|
|
assert sources["listed.safetensors"]["source"] == "model_list"
|
|
values["invalidate_local_hash_match_cache"].assert_called_once_with()
|
|
assert values["analyze_and_find_matches"].call_args.kwargs["analysis_id"] == (
|
|
"analysis-1"
|
|
)
|
|
values["self"]._update_analysis_progress.assert_any_call(
|
|
"analysis-1",
|
|
{"stage": "matching", "current": 1, "total": 5},
|
|
)
|
|
values["self"].analysis_progress.update.assert_called_once_with(
|
|
"analysis-1",
|
|
status="completed",
|
|
stage="completed",
|
|
message="Analysis complete",
|
|
current=4,
|
|
total=4,
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_analyze_route_returns_error_and_updates_progress_on_failure():
|
|
handlers, values = _build_routes(
|
|
{"analyze_and_find_matches": MagicMock(side_effect=RuntimeError("boom"))}
|
|
)
|
|
response = await handlers[("POST", "/model_resolver/analyze")](
|
|
_request({"workflow": {"nodes": []}, "analysis_id": "failed-analysis"})
|
|
)
|
|
|
|
assert response.status == 500
|
|
assert json.loads(response.text) == {"error": "boom"}
|
|
values["self"].analysis_progress.update.assert_called_once_with(
|
|
"failed-analysis",
|
|
status="error",
|
|
stage="error",
|
|
message="boom",
|
|
current=0,
|
|
total=0,
|
|
)
|
|
values["self"].logger.error.assert_called_once()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_workflow_resolution_route_returns_updated_workflow():
|
|
updated = {"nodes": [{"id": 1, "widgets_values": ["resolved.safetensors"]}]}
|
|
handlers, values = _build_routes(
|
|
{"apply_resolution": MagicMock(return_value=updated)}
|
|
)
|
|
workflow = {"nodes": []}
|
|
response = await handlers[("POST", "/model_resolver/resolve")](
|
|
_request(
|
|
{
|
|
"workflow": workflow,
|
|
"resolutions": [
|
|
{
|
|
"node_id": 1,
|
|
"widget_index": 0,
|
|
"resolved_path": "resolved.safetensors",
|
|
"category": "checkpoints",
|
|
}
|
|
],
|
|
}
|
|
)
|
|
)
|
|
|
|
assert response.status == 200
|
|
assert json.loads(response.text) == {"workflow": updated, "success": True}
|
|
values["apply_resolution"].assert_called_once_with(
|
|
workflow,
|
|
[
|
|
Resolution.from_mapping(
|
|
{
|
|
"node_id": 1,
|
|
"widget_index": 0,
|
|
"resolved_path": "resolved.safetensors",
|
|
"category": "checkpoints",
|
|
}
|
|
)
|
|
],
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
"resolution",
|
|
[
|
|
{
|
|
"node_id": 1,
|
|
"widget_index": "0",
|
|
"resolved_path": "resolved.safetensors",
|
|
},
|
|
{
|
|
"node_id": 1,
|
|
"widget_index": -1,
|
|
"resolved_path": "resolved.safetensors",
|
|
},
|
|
{
|
|
"node_id": 1,
|
|
"widget_index": 0,
|
|
"resolved_path": 123,
|
|
},
|
|
],
|
|
)
|
|
async def test_workflow_resolution_route_rejects_invalid_resolution_types(resolution):
|
|
handlers, values = _build_routes()
|
|
response = await handlers[("POST", "/model_resolver/resolve")](
|
|
_request(
|
|
{
|
|
"workflow": {"nodes": []},
|
|
"resolutions": [resolution],
|
|
}
|
|
)
|
|
)
|
|
|
|
assert response.status == 400
|
|
assert json.loads(response.text)["error"].startswith("Invalid resolution:")
|
|
values["apply_resolution"].assert_not_called()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_local_matches_route_rescans_and_passes_search_options():
|
|
matches = [
|
|
ModelMatch.from_mapping(
|
|
{
|
|
"model": {
|
|
"path": os.path.join("models", "model.safetensors")
|
|
}
|
|
}
|
|
)
|
|
]
|
|
handlers, values = _build_routes(
|
|
{"search_local_matches": MagicMock(return_value=matches)}
|
|
)
|
|
response = await handlers[("POST", "/model_resolver/local-matches")](
|
|
_request(
|
|
{
|
|
"filename": "model.safetensors",
|
|
"category": "loras",
|
|
"force_rescan": "true",
|
|
}
|
|
)
|
|
)
|
|
|
|
assert response.status == 200
|
|
assert json.loads(response.text) == {
|
|
"matches": [match.to_dict() for match in matches]
|
|
}
|
|
values["invalidate_local_hash_match_cache"].assert_called_once_with()
|
|
values["search_local_matches"].assert_called_once_with(
|
|
"model.safetensors",
|
|
category="loras",
|
|
similarity_threshold=0.0,
|
|
max_matches_per_model=10,
|
|
force_rescan=True,
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_local_matches_route_rejects_non_text_fields():
|
|
handlers, values = _build_routes()
|
|
response = await handlers[("POST", "/model_resolver/local-matches")](
|
|
_request({"filename": 123, "category": "checkpoints"})
|
|
)
|
|
|
|
assert response.status == 400
|
|
assert json.loads(response.text) == {
|
|
"error": "Local matches request filename must be a string"
|
|
}
|
|
values["search_local_matches"].assert_not_called()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_workflow_progress_route_delegates_analysis_tracker():
|
|
handlers, values = _build_routes()
|
|
request = SimpleNamespace()
|
|
response = await handlers[("GET", "/model_resolver/analyze-progress/{analysis_id}")](
|
|
request
|
|
)
|
|
|
|
assert response == {"status": "running"}
|
|
values["get_progress_response"].assert_called_once_with(
|
|
values["self"].analysis_progress,
|
|
request,
|
|
param_name="analysis_id",
|
|
not_found_payload={
|
|
"status": "unknown",
|
|
"stage": "unknown",
|
|
"message": "No analysis progress available",
|
|
"current": 0,
|
|
"total": 0,
|
|
},
|
|
)
|