Azornes-Comfyui-Model-Resolver/tests/test_workflow_service.py

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