kenyonxu-ComfyUI_KaiUtilities/tests/test_nodes.py

import json
import os
import sys
from unittest.mock import 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")
sys.path.insert(
    0, "/home/kai-remote/github/ComfyUI/.worktrees/kai-civitai-lora/custom_nodes"
)

from ComfyUI_KaiUtilities.nodes import KaiCivitaiLoRALoader


class TestKaiCivitaiLoRALoader:
    def test_inherits_lora_loader(self):
        from nodes import LoraLoader
        assert issubclass(KaiCivitaiLoRALoader, LoraLoader)

    def test_input_types_has_required_fields(self):
        inputs = KaiCivitaiLoRALoader.INPUT_TYPES()
        required = inputs.get("required", {})
        assert "model" in required
        assert "clip" in required
        assert "lora_name" in required
        assert "strength_model" in required
        assert "strength_clip" in required

    def test_return_types(self):
        assert KaiCivitaiLoRALoader.RETURN_TYPES == ("MODEL", "CLIP", "STRING", "STRING")

    def test_return_names(self):
        assert KaiCivitaiLoRALoader.RETURN_NAMES == (
            "model", "clip", "trigger_words", "civitai_json"
        )

    def test_function_name(self):
        assert KaiCivitaiLoRALoader.FUNCTION == "load_lora_with_info"

    def test_load_lora_with_info_returns_empty_when_no_cache(self):
        node = KaiCivitaiLoRALoader()
        node.load_lora = MagicMock(return_value=("model_out", "clip_out"))

        with patch("ComfyUI_KaiUtilities.nodes.cache_manager") as mock_cache:
            mock_cache.get_cache_for_file_sync = MagicMock(return_value=None)
            with patch("folder_paths.get_full_path_or_raise", return_value="/fake/lora.safetensors"):
                result = node.load_lora_with_info(
                    "model", "clip", "test_lora.safetensors", 1.0, 1.0
                )

        assert result == ("model_out", "clip_out", "", "")

    def test_load_lora_with_info_returns_trigger_words_and_json(self):
        node = KaiCivitaiLoRALoader()
        node.load_lora = MagicMock(return_value=("model_out", "clip_out"))

        cache_data = {
            "hash": "abc123",
            "name": "Test LoRA",
            "triggerWords": ["style", "masterpiece"],
            "fetched_at": "2026-04-20T10:00:00Z",
        }

        with patch("ComfyUI_KaiUtilities.nodes.cache_manager") as mock_cache:
            mock_cache.get_cache_for_file_sync = MagicMock(return_value=cache_data)
            with patch("folder_paths.get_full_path_or_raise", return_value="/fake/lora.safetensors"):
                result = node.load_lora_with_info(
                    "model", "clip", "test_lora.safetensors", 1.0, 1.0
                )

        assert result[0] == "model_out"
        assert result[1] == "clip_out"
        assert result[2] == "style, masterpiece"
        assert json.loads(result[3]) == cache_data

    def test_load_lora_with_info_path_error_returns_empty(self):
        node = KaiCivitaiLoRALoader()
        node.load_lora = MagicMock(return_value=("model_out", "clip_out"))

        with patch("folder_paths.get_full_path_or_raise", side_effect=FileNotFoundError("not found")):
            result = node.load_lora_with_info(
                "model", "clip", "missing.safetensors", 1.0, 1.0
            )

        assert result == ("model_out", "clip_out", "", "")

    def test_load_lora_with_info_trained_words_none(self):
        """trainedWords can be None from Civitai API."""
        node = KaiCivitaiLoRALoader()
        node.load_lora = MagicMock(return_value=("model_out", "clip_out"))

        cache_data = {
            "hash": "abc123",
            "name": "Test LoRA",
            "triggerWords": None,
            "fetched_at": "2026-04-20T10:00:00Z",
        }

        with patch("ComfyUI_KaiUtilities.nodes.cache_manager") as mock_cache:
            mock_cache.get_cache_for_file_sync = MagicMock(return_value=cache_data)
            with patch("folder_paths.get_full_path_or_raise", return_value="/fake/lora.safetensors"):
                result = node.load_lora_with_info(
                    "model", "clip", "test_lora.safetensors", 1.0, 1.0
                )

        assert result[2] == ""

    def test_load_lora_with_info_empty_trigger_words(self):
        """triggerWords can be empty list from Civitai API."""
        node = KaiCivitaiLoRALoader()
        node.load_lora = MagicMock(return_value=("model_out", "clip_out"))

        cache_data = {
            "hash": "abc123",
            "name": "Test LoRA",
            "triggerWords": [],
            "fetched_at": "2026-04-20T10:00:00Z",
        }

        with patch("ComfyUI_KaiUtilities.nodes.cache_manager") as mock_cache:
            mock_cache.get_cache_for_file_sync = MagicMock(return_value=cache_data)
            with patch("folder_paths.get_full_path_or_raise", return_value="/fake/lora.safetensors"):
                result = node.load_lora_with_info(
                    "model", "clip", "test_lora.safetensors", 1.0, 1.0
                )

        assert result[2] == ""

    def test_load_lora_with_info_nonlist_trigger_words(self):
        """triggerWords can be malformed (non-list)."""
        node = KaiCivitaiLoRALoader()
        node.load_lora = MagicMock(return_value=("model_out", "clip_out"))

        cache_data = {
            "hash": "abc123",
            "name": "Test LoRA",
            "triggerWords": "not a list",
            "fetched_at": "2026-04-20T10:00:00Z",
        }

        with patch("ComfyUI_KaiUtilities.nodes.cache_manager") as mock_cache:
            mock_cache.get_cache_for_file_sync = MagicMock(return_value=cache_data)
            with patch("folder_paths.get_full_path_or_raise", return_value="/fake/lora.safetensors"):
                result = node.load_lora_with_info(
                    "model", "clip", "test_lora.safetensors", 1.0, 1.0
                )

        assert result[2] == ""

    def test_load_lora_with_info_cache_lookup_error(self):
        """Cache lookup failure should not crash the node."""
        node = KaiCivitaiLoRALoader()
        node.load_lora = MagicMock(return_value=("model_out", "clip_out"))

        with patch("ComfyUI_KaiUtilities.nodes.cache_manager") as mock_cache:
            mock_cache.get_cache_for_file_sync = MagicMock(side_effect=RuntimeError("unexpected"))
            with patch("folder_paths.get_full_path_or_raise", return_value="/fake/lora.safetensors"):
                result = node.load_lora_with_info(
                    "model", "clip", "test_lora.safetensors", 1.0, 1.0
                )

        assert result == ("model_out", "clip_out", "", "")