diff --git a/kikotools/tools/embedding_autocomplete/node.py b/kikotools/tools/embedding_autocomplete/node.py index 05b129f..28edd02 100644 --- a/kikotools/tools/embedding_autocomplete/node.py +++ b/kikotools/tools/embedding_autocomplete/node.py @@ -5,7 +5,12 @@ Provides autocomplete functionality for embeddings and LoRAs in text inputs. import os from typing import Dict, List, Any -import folder_paths + +try: + import folder_paths +except ImportError: + # For testing outside ComfyUI environment + folder_paths = None class KikoEmbeddingAutocomplete: @@ -114,6 +119,8 @@ class KikoEmbeddingAutocomplete: # Get embedding files from ComfyUI's folder system try: print("[KikoEmbeddingAutocomplete] Getting embeddings list...") + if folder_paths is None: + return embeddings embedding_files = folder_paths.get_filename_list("embeddings") print( f"[KikoEmbeddingAutocomplete] Found {len(embedding_files)} embedding files" @@ -140,6 +147,8 @@ class KikoEmbeddingAutocomplete: # Get LoRA files from ComfyUI's folder system try: + if folder_paths is None: + return loras lora_files = folder_paths.get_filename_list("loras") for file in lora_files: name = os.path.splitext(file)[0] @@ -166,6 +175,8 @@ class KikoEmbeddingAutocomplete: # Check if embeddings/loras folders have changed try: + if folder_paths is None: + return 0 embeddings_path = folder_paths.get_folder_paths("embeddings")[0] loras_path = folder_paths.get_folder_paths("loras")[0] @@ -206,7 +217,10 @@ class KikoEmbeddingAutocompleteAPI: # Get embeddings if include_embeddings: try: - embedding_files = folder_paths.get_filename_list("embeddings") + if folder_paths is None: + embedding_files = [] + else: + embedding_files = folder_paths.get_filename_list("embeddings") for file in embedding_files: name = os.path.splitext(file)[0] match_name = name if case_sensitive else name.lower() @@ -238,7 +252,10 @@ class KikoEmbeddingAutocompleteAPI: # Get LoRAs if include_loras: try: - lora_files = folder_paths.get_filename_list("loras") + if folder_paths is None: + lora_files = [] + else: + lora_files = folder_paths.get_filename_list("loras") for file in lora_files: name = os.path.splitext(file)[0] match_name = name if case_sensitive else name.lower() diff --git a/tests/test_embedding_autocomplete.py b/tests/test_embedding_autocomplete.py index 5adff9d..0e76136 100644 --- a/tests/test_embedding_autocomplete.py +++ b/tests/test_embedding_autocomplete.py @@ -1,14 +1,7 @@ """Basic tests for KikoEmbeddingAutocomplete.""" import sys -from unittest.mock import MagicMock - -# Mock ComfyUI's folder_paths module BEFORE any imports -folder_paths_mock = MagicMock() -folder_paths_mock.get_filename_list = MagicMock(return_value=[]) -folder_paths_mock.get_folder_paths = MagicMock(return_value=["/mock/path"]) -folder_paths_mock.base_path = "/mock/base" -sys.modules["folder_paths"] = folder_paths_mock +from unittest.mock import MagicMock, patch def test_import(): @@ -58,34 +51,38 @@ def test_input_types(): def test_api_suggestions(): """Test the API suggestions method.""" - from kikotools.tools.embedding_autocomplete.node import KikoEmbeddingAutocompleteAPI + from kikotools.tools.embedding_autocomplete.node import ( + KikoEmbeddingAutocompleteAPI, + folder_paths, + ) - # Mock folder_paths to return some test files - sys.modules["folder_paths"].get_filename_list = MagicMock( - side_effect=lambda x: ( - ["test1.pt", "test2.safetensors"] - if x == "embeddings" - else ["lora1.pt", "lora2.safetensors"] + # Mock folder_paths if it exists (will be None in tests) + with patch("kikotools.tools.embedding_autocomplete.node.folder_paths") as mock_fp: + mock_fp.get_filename_list = MagicMock( + side_effect=lambda x: ( + ["test1.pt", "test2.safetensors"] + if x == "embeddings" + else ["lora1.pt", "lora2.safetensors"] + ) ) - ) - # Test with embeddings - suggestions = KikoEmbeddingAutocompleteAPI.get_suggestions( - prefix="test", include_embeddings=True, include_loras=False - ) + # Test with embeddings + suggestions = KikoEmbeddingAutocompleteAPI.get_suggestions( + prefix="test", include_embeddings=True, include_loras=False + ) - assert len(suggestions) == 2 - assert suggestions[0]["type"] == "embedding" - assert suggestions[0]["name"] == "test1" + assert len(suggestions) == 2 + assert suggestions[0]["type"] == "embedding" + assert suggestions[0]["name"] == "test1" - # Test with LoRAs - suggestions = KikoEmbeddingAutocompleteAPI.get_suggestions( - prefix="lora", include_embeddings=False, include_loras=True - ) + # Test with LoRAs + suggestions = KikoEmbeddingAutocompleteAPI.get_suggestions( + prefix="lora", include_embeddings=False, include_loras=True + ) - assert len(suggestions) == 2 - assert suggestions[0]["type"] == "lora" - assert "