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