ComfyUI-Lora-Manager/tests/routes/test_download_routing_handlers.py

"""Tests for the download routing HTTP handler."""

import json

import pytest

from py.routes.handlers.download_routing_handlers import DownloadRoutingHandler


class FakeRequest:
    def __init__(self, payload):
        self._payload = payload

    async def json(self):
        if isinstance(self._payload, Exception):
            raise self._payload
        return self._payload


@pytest.mark.asyncio
async def test_diffusion_base_model_routes_to_unet():
    """The reported Anima case: file type "Model", baseModel "Anima"."""
    handler = DownloadRoutingHandler()
    response = await handler.get_download_routing(
        FakeRequest(
            {"model_type": "checkpoint", "base_model": "Anima", "file_types": ["Model"]}
        )
    )
    payload = json.loads(response.text)
    assert response.status == 200
    assert payload == {"success": True, "is_diffusion_model": True, "root_kind": "unet"}


@pytest.mark.asyncio
async def test_unet_file_type_routes_to_unet():
    handler = DownloadRoutingHandler()
    response = await handler.get_download_routing(
        FakeRequest(
            {"model_type": "checkpoint", "base_model": "SDXL 1.0", "file_types": ["UNet"]}
        )
    )
    payload = json.loads(response.text)
    assert payload["is_diffusion_model"] is True
    assert payload["root_kind"] == "unet"


@pytest.mark.asyncio
async def test_regular_checkpoint_stays_on_checkpoint_root():
    handler = DownloadRoutingHandler()
    response = await handler.get_download_routing(
        FakeRequest(
            {"model_type": "checkpoint", "base_model": "SDXL 1.0", "file_types": ["Model"]}
        )
    )
    payload = json.loads(response.text)
    assert payload["is_diffusion_model"] is False
    assert payload["root_kind"] == "checkpoint"


@pytest.mark.asyncio
async def test_lora_is_never_diffusion():
    handler = DownloadRoutingHandler()
    response = await handler.get_download_routing(
        FakeRequest({"model_type": "lora", "base_model": "Anima", "file_types": []})
    )
    payload = json.loads(response.text)
    assert payload["is_diffusion_model"] is False
    assert payload["root_kind"] == "lora"


@pytest.mark.asyncio
async def test_missing_model_type_rejected():
    handler = DownloadRoutingHandler()
    response = await handler.get_download_routing(FakeRequest({"base_model": "Anima"}))
    assert response.status == 400


@pytest.mark.asyncio
async def test_invalid_file_types_rejected():
    handler = DownloadRoutingHandler()
    response = await handler.get_download_routing(
        FakeRequest({"model_type": "checkpoint", "file_types": "Model"})
    )
    assert response.status == 400


@pytest.mark.asyncio
async def test_invalid_json_rejected():
    handler = DownloadRoutingHandler()
    response = await handler.get_download_routing(
        FakeRequest(json.JSONDecodeError("bad", "", 0))
    )
    assert response.status == 400


@pytest.mark.asyncio
async def test_other_model_type_returns_sub_type():
    handler = DownloadRoutingHandler()
    response = await handler.get_download_routing(
        FakeRequest({"model_type": "TextEncoder", "file_types": ["Model"]})
    )
    payload = json.loads(response.text)
    assert response.status == 200
    assert payload == {"success": True, "root_kind": "other", "sub_type": "text_encoder"}


@pytest.mark.asyncio
async def test_other_explicit_file_pick_wins():
    handler = DownloadRoutingHandler()
    response = await handler.get_download_routing(
        FakeRequest(
            {
                "model_type": "Other",
                "file_types": ["Model"],
                "selected_file_type": "VAE",
            }
        )
    )
    payload = json.loads(response.text)
    assert payload["root_kind"] == "other"
    assert payload["sub_type"] == "vae"


@pytest.mark.asyncio
async def test_other_file_type_fallback_when_model_type_unmapped():
    handler = DownloadRoutingHandler()
    response = await handler.get_download_routing(
        FakeRequest({"model_type": "Other", "file_types": ["Model", "Upscaler"]})
    )
    payload = json.loads(response.text)
    assert payload["sub_type"] == "upscaler"


@pytest.mark.asyncio
async def test_other_undecidable_sub_type_is_none():
    handler = DownloadRoutingHandler()
    response = await handler.get_download_routing(
        FakeRequest({"model_type": "Other", "file_types": ["Model"]})
    )
    payload = json.loads(response.text)
    assert response.status == 200
    assert payload == {"success": True, "root_kind": "other", "sub_type": None}


@pytest.mark.asyncio
async def test_other_invalid_selected_file_type_rejected():
    handler = DownloadRoutingHandler()
    response = await handler.get_download_routing(
        FakeRequest({"model_type": "VAE", "selected_file_type": 123})
    )
    assert response.status == 400