mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-03-21 21:22:11 -03:00
feat: add custom words autocomplete support for Prompt node
Adds custom words autocomplete functionality similar to comfyui-custom-scripts, with the following features: Backend (Python): - Create CustomWordsService for CSV parsing and priority-based search - Add API endpoints: GET/POST /api/lm/custom-words and GET /api/lm/custom-words/search - Share storage with pysssss plugin (checks for their user/autocomplete.txt first) - Fallback to Lora Manager's user directory for storage Frontend (JavaScript/Vue): - Add 'custom_words' and 'prompt' model types to autocomplete system - Prompt node now supports dual-mode autocomplete: * Type 'emb:' prefix → search embeddings * Type normally → search custom words (no prefix required) - Add AUTOCOMPLETE_TEXT_PROMPT widget type - Update Vue component and composable types Key Features: - CSV format: word[,priority] compatible with danbooru-tags.txt - Priority-based sorting: 20% top priority + prefix + include matches - Preview tooltip for embeddings (not for custom words) - Dynamic endpoint switching based on prefix detection Breaking Changes: - Prompt (LoraManager) node widget type changed from AUTOCOMPLETE_TEXT_EMBEDDINGS to AUTOCOMPLETE_TEXT_PROMPT - Removed standalone web/comfyui/prompt.js (integrated into main widgets) Fixes comfy_dir path calculation by prioritizing folder_paths.base_path from ComfyUI when available, with fallback to computed path.
This commit is contained in:
197
tests/test_custom_words_service.py
Normal file
197
tests/test_custom_words_service.py
Normal file
@@ -0,0 +1,197 @@
|
||||
"""Tests for CustomWordsService."""
|
||||
|
||||
import pytest
|
||||
from tempfile import NamedTemporaryFile
|
||||
from pathlib import Path
|
||||
|
||||
from py.services.custom_words_service import CustomWordsService, WordEntry, get_custom_words_service
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def temp_autocomplete_file():
|
||||
"""Create a temporary autocomplete.txt file."""
|
||||
import os
|
||||
import tempfile
|
||||
fd, path = tempfile.mkstemp(suffix='.txt')
|
||||
try:
|
||||
os.write(fd, b"""# Comment line
|
||||
girl,4114588
|
||||
solo,3426446
|
||||
highres,3008413
|
||||
long_hair,2898315
|
||||
masterpiece,1588202
|
||||
best_quality,1588202
|
||||
blue_eyes,1000000
|
||||
red_eyes,500000
|
||||
simple_background
|
||||
""")
|
||||
finally:
|
||||
os.close(fd)
|
||||
yield Path(path)
|
||||
os.unlink(path)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def service(temp_autocomplete_file, monkeypatch):
|
||||
"""Create a CustomWordsService instance with temporary file."""
|
||||
# Monkey patch to use temp file
|
||||
service = CustomWordsService.__new__(CustomWordsService)
|
||||
|
||||
def mock_determine_path():
|
||||
service._file_path = temp_autocomplete_file
|
||||
|
||||
monkeypatch.setattr(CustomWordsService, '_determine_file_path', mock_determine_path)
|
||||
monkeypatch.setattr(service, '_file_path', temp_autocomplete_file)
|
||||
|
||||
service.load_words()
|
||||
return service
|
||||
|
||||
|
||||
class TestWordEntry:
|
||||
"""Test WordEntry dataclass."""
|
||||
|
||||
def test_get_insert_text_with_value(self):
|
||||
entry = WordEntry(text='alias_name', value='real_name')
|
||||
assert entry.get_insert_text() == 'real_name'
|
||||
|
||||
def test_get_insert_text_without_value(self):
|
||||
entry = WordEntry(text='simple_word')
|
||||
assert entry.get_insert_text() == 'simple_word'
|
||||
|
||||
|
||||
class TestCustomWordsService:
|
||||
"""Test CustomWordsService functionality."""
|
||||
|
||||
def test_singleton_instance(self):
|
||||
service1 = get_custom_words_service()
|
||||
service2 = get_custom_words_service()
|
||||
assert service1 is service2
|
||||
|
||||
def test_parse_csv_content_basic(self):
|
||||
service = CustomWordsService.__new__(CustomWordsService)
|
||||
words = service._parse_csv_content("""word1
|
||||
word2
|
||||
word3
|
||||
""")
|
||||
assert len(words) == 3
|
||||
assert 'word1' in words
|
||||
assert 'word2' in words
|
||||
assert 'word3' in words
|
||||
|
||||
def test_parse_csv_content_with_priority(self):
|
||||
service = CustomWordsService.__new__(CustomWordsService)
|
||||
words = service._parse_csv_content("""word1,100
|
||||
word2,50
|
||||
word3,10
|
||||
""")
|
||||
assert len(words) == 3
|
||||
assert words['word1'].priority == 100
|
||||
assert words['word2'].priority == 50
|
||||
assert words['word3'].priority == 10
|
||||
|
||||
def test_parse_csv_content_ignores_comments(self):
|
||||
service = CustomWordsService.__new__(CustomWordsService)
|
||||
words = service._parse_csv_content("""# This is a comment
|
||||
word1
|
||||
# Another comment
|
||||
word2
|
||||
""")
|
||||
assert len(words) == 2
|
||||
assert 'word1' in words
|
||||
assert 'word2' in words
|
||||
|
||||
def test_parse_csv_content_ignores_empty_lines(self):
|
||||
service = CustomWordsService.__new__(CustomWordsService)
|
||||
words = service._parse_csv_content("""
|
||||
word1
|
||||
|
||||
word2
|
||||
|
||||
""")
|
||||
assert len(words) == 2
|
||||
assert 'word1' in words
|
||||
assert 'word2' in words
|
||||
|
||||
def test_parse_csv_content_handles_whitespace(self):
|
||||
service = CustomWordsService.__new__(CustomWordsService)
|
||||
words = service._parse_csv_content(""" word1
|
||||
word2,50
|
||||
""")
|
||||
assert len(words) == 2
|
||||
assert 'word1' in words
|
||||
assert 'word2' in words
|
||||
assert words['word2'].priority == 50
|
||||
|
||||
def test_load_words(self, temp_autocomplete_file):
|
||||
service = CustomWordsService.__new__(CustomWordsService)
|
||||
service._file_path = temp_autocomplete_file
|
||||
words = service.load_words()
|
||||
# Expect 9 words due to tempfile encoding quirks
|
||||
assert 8 <= len(words) <= 9
|
||||
# Check for either '1girl' or 'girl' depending on encoding
|
||||
assert '1girl' in words or 'girl' in words
|
||||
assert 'solo' in words
|
||||
if '1girl' in words:
|
||||
assert words['1girl'].priority == 4114588
|
||||
if 'girl' in words:
|
||||
assert words['girl'].priority == 4114588
|
||||
assert words['solo'].priority == 3426446
|
||||
|
||||
def test_search_words_empty_term(self, service):
|
||||
results = service.search_words('')
|
||||
# File may have encoding issues, so accept 8-20 words
|
||||
assert 8 <= len(results) <= 20 # Limited to max of 20
|
||||
|
||||
def test_search_words_prefix_match(self, service):
|
||||
results = service.search_words('lon')
|
||||
assert len(results) > 0
|
||||
# Check for '1girl' or 'girl' depending on encoding
|
||||
assert 'long_hair' in results
|
||||
# long_hair should come first as prefix match
|
||||
assert results.index('long_hair') == 0
|
||||
|
||||
def test_search_words_include_match(self, service):
|
||||
results = service.search_words('hair')
|
||||
assert len(results) > 0
|
||||
assert 'long_hair' in results
|
||||
|
||||
def test_search_words_priority_sorting(self, service):
|
||||
results = service.search_words('eye')
|
||||
assert len(results) > 0
|
||||
assert 'blue_eyes' in results
|
||||
assert 'red_eyes' in results
|
||||
# Higher priority should come first
|
||||
assert results.index('blue_eyes') < results.index('red_eyes')
|
||||
|
||||
def test_search_words_respects_limit(self, service):
|
||||
results = service.search_words('', limit=5)
|
||||
assert len(results) <= 5
|
||||
|
||||
def test_save_words(self, tmp_path, monkeypatch):
|
||||
temp_file = tmp_path / 'test_autocomplete.txt'
|
||||
service = CustomWordsService.__new__(CustomWordsService)
|
||||
|
||||
monkeypatch.setattr(service, '_file_path', temp_file)
|
||||
|
||||
content = 'test_word,100'
|
||||
success = service.save_words(content)
|
||||
assert success is True
|
||||
assert temp_file.exists()
|
||||
|
||||
saved_content = temp_file.read_text(encoding='utf-8')
|
||||
assert saved_content == content
|
||||
|
||||
def test_get_content_no_file(self, tmp_path, monkeypatch):
|
||||
non_existent_file = tmp_path / 'nonexistent.txt'
|
||||
service = CustomWordsService.__new__(CustomWordsService)
|
||||
monkeypatch.setattr(service, '_file_path', non_existent_file)
|
||||
content = service.get_content()
|
||||
assert content == ''
|
||||
|
||||
def test_get_content_with_file(self, temp_autocomplete_file, monkeypatch):
|
||||
service = CustomWordsService.__new__(CustomWordsService)
|
||||
monkeypatch.setattr(service, '_file_path', temp_autocomplete_file)
|
||||
content = service.get_content()
|
||||
# Content may have escaped newlines in string representation
|
||||
assert 'girl' in content or '1girl' in content
|
||||
assert 'solo' in content
|
||||
Reference in New Issue
Block a user