kenyonxu-ComfyUI_KaiUtilities/tests/test_init.py

import asyncio
import json
import os
import sys
from unittest.mock import AsyncMock, MagicMock, patch

import pytest

sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
sys.path.insert(0, "/home/kai-remote/github/ComfyUI/.worktrees/kai-civitai-lora")

from ComfyUI_KaiUtilities import _normalize_civitai_response, fetch_civitai_info, get_cache, get_image


def _run_async(coro):
    return asyncio.run(coro)


class TestGetImage:
    def test_get_image_invalid_hash_format(self):
        mock_request = MagicMock()
        mock_request.match_info = {"hash": "../../../etc/passwd", "index": "0"}
        response = _run_async(get_image(mock_request))
        assert response.status == 400

    def test_get_image_invalid_hash_length(self):
        mock_request = MagicMock()
        mock_request.match_info = {"hash": "abc123", "index": "0"}
        response = _run_async(get_image(mock_request))
        assert response.status == 400

    def test_get_image_invalid_index(self):
        mock_request = MagicMock()
        mock_request.match_info = {"hash": "a" * 64, "index": "notanumber"}
        response = _run_async(get_image(mock_request))
        assert response.status == 400


class TestNormalizeCivitaiResponse:
    def test_basic_fields(self):
        raw = {
            "modelId": 12345,
            "name": "Test Model",
            "baseModel": "SDXL",
            "trainedWords": ["style", "masterpiece"],
            "model": {"description": "A test model"},
            "images": [
                {"url": "https://civitai.com/img1.jpg", "type": "image", "nsfw": True},
                {"url": "https://civitai.com/img2.jpg", "type": "image", "nsfw": False},
            ],
            "files": [
                {"hashes": {"SHA256": "ABC123"}, "downloadUrl": "https://civitai.com/download"}
            ],
        }
        result = _normalize_civitai_response(raw, "abc123")
        assert result["hash"] == "abc123"
        assert result["modelId"] == 12345
        assert result["name"] == "Test Model"
        assert result["baseModel"] == "SDXL"
        assert result["triggerWords"] == ["style", "masterpiece"]
        assert result["description"] == "A test model"
        assert len(result["images"]) == 2
        assert result["images"][0]["nsfw"] is True
        assert result["images"][1]["nsfw"] is False
        assert result["downloadUrl"] == "https://civitai.com/download"

    def test_trained_words_none(self):
        raw = {
            "name": "No Words",
            "trainedWords": None,
            "model": {},
            "images": [],
            "files": [],
        }
        result = _normalize_civitai_response(raw, "hash123")
        assert result["triggerWords"] == []

    def test_model_is_none(self):
        raw = {
            "name": "No Model",
            "trainedWords": [],
            "model": None,
            "images": [],
            "files": [],
        }
        result = _normalize_civitai_response(raw, "hash123")
        assert result["description"] == ""

    def test_nsfw_uses_boolean_field(self):
        """nsfw is a boolean on image objects, not nsfwLevel."""
        raw = {
            "name": "NSFW Test",
            "trainedWords": [],
            "model": {},
            "images": [
                {"url": "https://civitai.com/img.jpg", "nsfw": True},
            ],
            "files": [],
        }
        result = _normalize_civitai_response(raw, "hash123")
        assert result["images"][0]["nsfw"] is True

    def test_sha256_lowercase_fallback(self):
        raw = {
            "name": "Case Test",
            "trainedWords": [],
            "model": {},
            "images": [],
            "files": [
                {"hashes": {"sha256": "def456"}, "downloadUrl": "https://lowercase.com/dl"}
            ],
        }
        result = _normalize_civitai_response(raw, "def456")
        assert result["downloadUrl"] == "https://lowercase.com/dl"


class TestFetchCivitaiInfo:
    def test_missing_lora_name(self):
        mock_request = MagicMock()
        mock_request.json = AsyncMock(return_value={})
        response = _run_async(fetch_civitai_info(mock_request))
        assert response.status == 400
        body = json.loads(response.text)
        assert "Missing lora_name" in body["error"]

    def test_invalid_lora_name(self):
        mock_request = MagicMock()
        mock_request.json = AsyncMock(return_value={"lora_name": "nonexistent.safetensors"})
        with patch("folder_paths.get_filename_list", return_value=["valid_lora.safetensors"]):
            response = _run_async(fetch_civitai_info(mock_request))
        assert response.status == 400
        body = json.loads(response.text)
        assert "Invalid lora_name" in body["error"]

    def test_fetch_civitai_info_hash_failure(self):
        mock_request = MagicMock()
        mock_request.json = AsyncMock(return_value={"lora_name": "test.safetensors"})
        with patch("folder_paths.get_filename_list", return_value=["test.safetensors"]):
            with patch("folder_paths.get_full_path_or_raise", return_value="/fake/test.safetensors"):
                with patch("ComfyUI_KaiUtilities.compute_file_hash", side_effect=RuntimeError("hash failed")):
                    response = _run_async(fetch_civitai_info(mock_request))
        assert response.status == 500
        body = json.loads(response.text)
        assert "Hash failed" in body["error"]


class TestGetCache:
    def test_get_cache_missing(self):
        mock_request = MagicMock()
        mock_request.match_info = {"hash": "a" * 64}
        with patch("ComfyUI_KaiUtilities.cache_manager") as mock_cache:
            mock_cache.read_cache = MagicMock(return_value=None)
            response = _run_async(get_cache(mock_request))
        assert response.status == 404
        body = json.loads(response.text)
        assert body["exists"] is False

    def test_get_cache_exists(self):
        mock_request = MagicMock()
        mock_request.match_info = {"hash": "a" * 64}
        cache_data = {"hash": "a" * 64, "name": "Test"}
        with patch("ComfyUI_KaiUtilities.cache_manager") as mock_cache:
            mock_cache.read_cache = MagicMock(return_value=cache_data)
            response = _run_async(get_cache(mock_request))
        assert response.status == 200
        body = json.loads(response.text)
        assert body["exists"] is True
        assert body["data"] == cache_data