Vito
8146250274
* fix(search): PromptSearchList crash, filters, and batch output (#52) - Fix empty list crash (OUTPUT_IS_LIST requires at least one element) - Add partial tag matching (LIKE instead of exact) - Add preview and count outputs for batch visibility - Add skip_multipart filter for Clip_N video prompts - Add LoRA-only prompt filter - Collapse newlines for StringOutputList compatibility * feat(db): add tag_partial parameter for LIKE tag matching (#52) Allows search_prompts to use LIKE matching for tags instead of exact. Used by PromptSearchList for discovery-style search. Existing callers default to exact matching (unchanged behavior). * feat(tracker): add FIFO prompt queue for batch image linking (#52) Batch workflows encode all prompts first, then save all images later. The queue preserves encode order so each image links to the correct prompt by position (pop on save matches push on encode). * fix(monitor): multi-strategy image-to-prompt linking (#52) Previous approach relied on live prompt tracker state, which was always stale for batch workflows (all images linked to last prompt). New 5-strategy priority chain: 1. FIFO queue pop (batch-correct by position) 2. Metadata-based DB lookup (extracts prompt from PNG) 3. Prompt snapshot at file-creation time 4. Live prompt tracker 5. Most recent DB prompt fallback * test: add tests for prompt queue and search filters (#52) - 6 queue tests: FIFO ordering, batch simulation, exhaustion - 15 filter tests: newline collapse, Clip_ filter, LoRA-only filter - Fix _make_tracker helper to include queue fields * chore: bump version to 3.2.2 * fix: address code review — regex backtracking and queue expiry (#52) - Replace backtracking-prone LoRA-only regex with subtraction approach (re.sub removes tags, check if anything remains). Fixes CodeQL alert. - Add expiry check to pop_next_prompt — skips stale entries to prevent incorrect linking from accumulated queue entries across sessions. - Cleanup thread now also prunes expired queue entries.
317 lines
12 KiB
Python
317 lines
12 KiB
Python
"""
|
|
Tests for prompt tracking system.
|
|
|
|
Tests PromptTracker, PromptExecutionContext, and the singleton
|
|
get_prompt_tracker. Uses a mock db_manager to avoid database deps.
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
import time
|
|
import threading
|
|
import unittest
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
|
|
from utils.prompt_tracker import PromptTracker, PromptExecutionContext
|
|
|
|
|
|
def _make_tracker(prompt_timeout=600, cleanup_interval=300):
|
|
"""Create a PromptTracker with mocked db_manager and config."""
|
|
mock_db = MagicMock()
|
|
mock_db.get_prompt_by_hash.return_value = None
|
|
|
|
# Patch GalleryConfig import to avoid needing ComfyUI server
|
|
with patch.object(PromptTracker, "__init__", lambda self, db: None):
|
|
tracker = PromptTracker.__new__(PromptTracker)
|
|
|
|
# Manual init to avoid the config import and cleanup thread
|
|
from utils.logging_config import get_logger
|
|
|
|
tracker.logger = get_logger("prompt_manager.prompt_tracker")
|
|
tracker.db_manager = mock_db
|
|
tracker._local = threading.local()
|
|
tracker.active_prompts = {}
|
|
tracker.lock = threading.Lock()
|
|
tracker._prompt_queue = []
|
|
tracker._queue_lock = threading.Lock()
|
|
tracker.cleanup_interval = cleanup_interval
|
|
tracker.prompt_timeout = prompt_timeout
|
|
# Don't start cleanup thread in tests
|
|
return tracker
|
|
|
|
|
|
class TestPromptTrackerBasic(unittest.TestCase):
|
|
"""Test basic PromptTracker operations."""
|
|
|
|
def setUp(self):
|
|
self.tracker = _make_tracker()
|
|
|
|
def test_set_current_prompt_returns_execution_id(self):
|
|
exec_id = self.tracker.set_current_prompt("test prompt", {"prompt_id": 42})
|
|
self.assertTrue(exec_id.startswith("exec_"))
|
|
|
|
def test_set_current_prompt_stores_in_active(self):
|
|
exec_id = self.tracker.set_current_prompt("test", {"prompt_id": 1})
|
|
self.assertIn(exec_id, self.tracker.active_prompts)
|
|
|
|
def test_get_current_prompt_returns_set_prompt(self):
|
|
self.tracker.set_current_prompt("hello world", {"prompt_id": 5})
|
|
current = self.tracker.get_current_prompt()
|
|
self.assertIsNotNone(current)
|
|
self.assertEqual(current["text"], "hello world")
|
|
self.assertEqual(current["id"], 5)
|
|
|
|
def test_get_current_prompt_returns_none_when_empty(self):
|
|
result = self.tracker.get_current_prompt()
|
|
self.assertIsNone(result)
|
|
|
|
def test_clear_current_prompt(self):
|
|
self.tracker.set_current_prompt("to clear", {"prompt_id": 1})
|
|
self.tracker.clear_current_prompt()
|
|
result = self.tracker.get_current_prompt()
|
|
self.assertIsNone(result)
|
|
|
|
def test_clear_removes_from_active(self):
|
|
exec_id = self.tracker.set_current_prompt("test", {"prompt_id": 1})
|
|
self.tracker.clear_current_prompt()
|
|
self.assertNotIn(exec_id, self.tracker.active_prompts)
|
|
|
|
def test_clear_when_no_prompt_does_not_raise(self):
|
|
# Should not raise
|
|
self.tracker.clear_current_prompt()
|
|
|
|
|
|
class TestPromptTrackerTimeout(unittest.TestCase):
|
|
"""Test timeout behavior."""
|
|
|
|
def test_expired_prompt_returns_none(self):
|
|
tracker = _make_tracker(prompt_timeout=0.1)
|
|
tracker.set_current_prompt("expires fast", {"prompt_id": 1})
|
|
time.sleep(0.2)
|
|
result = tracker.get_current_prompt()
|
|
self.assertIsNone(result)
|
|
|
|
def test_extend_timeout_resets_timestamp(self):
|
|
tracker = _make_tracker(prompt_timeout=1)
|
|
exec_id = tracker.set_current_prompt("test", {"prompt_id": 1})
|
|
old_ts = tracker.active_prompts[exec_id]["timestamp"]
|
|
time.sleep(0.1)
|
|
tracker.extend_prompt_timeout(exec_id)
|
|
new_ts = tracker.active_prompts[exec_id]["timestamp"]
|
|
self.assertGreater(new_ts, old_ts)
|
|
|
|
def test_extend_nonexistent_id_does_not_raise(self):
|
|
tracker = _make_tracker()
|
|
tracker.extend_prompt_timeout("nonexistent_id")
|
|
|
|
|
|
class TestPromptTrackerMultiplePrompts(unittest.TestCase):
|
|
"""Test multiple prompt tracking."""
|
|
|
|
def setUp(self):
|
|
self.tracker = _make_tracker()
|
|
|
|
def test_get_active_prompts(self):
|
|
self.tracker.set_current_prompt("p1", {"prompt_id": 1})
|
|
self.tracker.set_current_prompt("p2", {"prompt_id": 2})
|
|
active = self.tracker.get_active_prompts()
|
|
self.assertEqual(len(active), 2)
|
|
|
|
def test_get_active_prompts_returns_copy(self):
|
|
self.tracker.set_current_prompt("p1", {"prompt_id": 1})
|
|
active = self.tracker.get_active_prompts()
|
|
active.clear()
|
|
# Original should be untouched
|
|
self.assertEqual(len(self.tracker.active_prompts), 1)
|
|
|
|
def test_clear_all_active_prompts(self):
|
|
self.tracker.set_current_prompt("p1", {"prompt_id": 1})
|
|
self.tracker.set_current_prompt("p2", {"prompt_id": 2})
|
|
cleared = self.tracker.clear_all_active_prompts()
|
|
self.assertEqual(cleared, 2)
|
|
self.assertEqual(len(self.tracker.active_prompts), 0)
|
|
|
|
def test_clear_all_returns_zero_when_empty(self):
|
|
cleared = self.tracker.clear_all_active_prompts()
|
|
self.assertEqual(cleared, 0)
|
|
|
|
|
|
class TestPromptTrackerFallback(unittest.TestCase):
|
|
"""Test fallback prompt lookup (for cross-thread access)."""
|
|
|
|
def test_fallback_finds_recent_prompt(self):
|
|
tracker = _make_tracker()
|
|
# Set prompt in active_prompts directly (simulating another thread)
|
|
tracker.active_prompts["exec_123"] = {
|
|
"id": 42,
|
|
"execution_id": "exec_123",
|
|
"text": "from other thread",
|
|
"timestamp": time.time(),
|
|
"thread_id": 0,
|
|
"additional_data": {},
|
|
}
|
|
# Current thread has no local prompt, but fallback should find it
|
|
result = tracker.get_current_prompt()
|
|
self.assertIsNotNone(result)
|
|
self.assertEqual(result["text"], "from other thread")
|
|
|
|
def test_fallback_uses_db_hash_lookup(self):
|
|
tracker = _make_tracker()
|
|
tracker.db_manager.get_prompt_by_hash.return_value = {"id": 99, "text": "found"}
|
|
exec_id = tracker.set_current_prompt("lookup test")
|
|
self.assertEqual(tracker.active_prompts[exec_id]["id"], 99)
|
|
|
|
def test_fallback_generates_temp_id(self):
|
|
tracker = _make_tracker()
|
|
tracker.db_manager.get_prompt_by_hash.return_value = None
|
|
exec_id = tracker.set_current_prompt("not in db")
|
|
prompt_id = tracker.active_prompts[exec_id]["id"]
|
|
self.assertTrue(str(prompt_id).startswith("temp_"))
|
|
|
|
|
|
class TestPromptTrackerStatus(unittest.TestCase):
|
|
"""Test get_status method."""
|
|
|
|
def test_status_empty_tracker(self):
|
|
tracker = _make_tracker()
|
|
status = tracker.get_status()
|
|
self.assertEqual(status["active_prompts_count"], 0)
|
|
self.assertIsNone(status["current_prompt_id"])
|
|
self.assertIsNone(status["current_execution_id"])
|
|
self.assertEqual(status["prompt_timeout"], 600)
|
|
self.assertEqual(status["cleanup_interval"], 300)
|
|
|
|
def test_status_with_active_prompt(self):
|
|
tracker = _make_tracker()
|
|
tracker.set_current_prompt("active", {"prompt_id": 7})
|
|
status = tracker.get_status()
|
|
self.assertEqual(status["active_prompts_count"], 1)
|
|
self.assertEqual(status["current_prompt_id"], 7)
|
|
self.assertIsNotNone(status["current_execution_id"])
|
|
|
|
|
|
class TestGenerateExecutionId(unittest.TestCase):
|
|
"""Test execution ID generation."""
|
|
|
|
def test_format(self):
|
|
tracker = _make_tracker()
|
|
exec_id = tracker.generate_execution_id()
|
|
self.assertTrue(exec_id.startswith("exec_"))
|
|
parts = exec_id.split("_")
|
|
self.assertEqual(len(parts), 3) # exec, uuid8, timestamp
|
|
|
|
def test_unique(self):
|
|
tracker = _make_tracker()
|
|
ids = {tracker.generate_execution_id() for _ in range(100)}
|
|
self.assertEqual(len(ids), 100)
|
|
|
|
|
|
class TestPromptExecutionContext(unittest.TestCase):
|
|
"""Test the context manager."""
|
|
|
|
def test_context_manager_sets_prompt(self):
|
|
tracker = _make_tracker()
|
|
with PromptExecutionContext(tracker, "context test", prompt_id=10) as exec_id:
|
|
self.assertIsNotNone(exec_id)
|
|
current = tracker.get_current_prompt()
|
|
self.assertEqual(current["text"], "context test")
|
|
|
|
def test_context_manager_does_not_clear_on_exit(self):
|
|
tracker = _make_tracker()
|
|
with PromptExecutionContext(tracker, "stays active", prompt_id=1):
|
|
pass
|
|
# Prompt should still be in active_prompts (designed behavior)
|
|
self.assertEqual(len(tracker.active_prompts), 1)
|
|
|
|
def test_context_manager_with_exception(self):
|
|
tracker = _make_tracker()
|
|
try:
|
|
with PromptExecutionContext(tracker, "error test", prompt_id=1):
|
|
raise ValueError("test error")
|
|
except ValueError:
|
|
pass
|
|
# Should not crash, prompt stays tracked
|
|
self.assertEqual(len(tracker.active_prompts), 1)
|
|
|
|
|
|
class TestCleanupExpiredPrompts(unittest.TestCase):
|
|
"""Test the cleanup logic (not the thread, just the logic)."""
|
|
|
|
def test_expired_prompts_cleaned(self):
|
|
tracker = _make_tracker(prompt_timeout=0.1)
|
|
tracker.set_current_prompt("old", {"prompt_id": 1})
|
|
time.sleep(0.2)
|
|
|
|
# Manually run cleanup logic
|
|
current_time = time.time()
|
|
expired_ids = []
|
|
with tracker.lock:
|
|
for exec_id, prompt_data in tracker.active_prompts.items():
|
|
if current_time - prompt_data["timestamp"] > tracker.prompt_timeout:
|
|
expired_ids.append(exec_id)
|
|
for exec_id in expired_ids:
|
|
tracker.active_prompts.pop(exec_id, None)
|
|
|
|
self.assertEqual(len(tracker.active_prompts), 0)
|
|
self.assertEqual(len(expired_ids), 1)
|
|
|
|
|
|
class TestPromptQueue(unittest.TestCase):
|
|
"""Test the FIFO prompt queue for batch workflow image linking."""
|
|
|
|
def setUp(self):
|
|
self.tracker = _make_tracker()
|
|
|
|
def test_queue_empty_initially(self):
|
|
result = self.tracker.pop_next_prompt()
|
|
self.assertIsNone(result)
|
|
|
|
def test_set_prompt_pushes_to_queue(self):
|
|
self.tracker.set_current_prompt("p1", {"prompt_id": 10})
|
|
self.assertEqual(len(self.tracker._prompt_queue), 1)
|
|
|
|
def test_pop_returns_fifo_order(self):
|
|
self.tracker.set_current_prompt("first", {"prompt_id": 1})
|
|
self.tracker.set_current_prompt("second", {"prompt_id": 2})
|
|
self.tracker.set_current_prompt("third", {"prompt_id": 3})
|
|
|
|
p1 = self.tracker.pop_next_prompt()
|
|
p2 = self.tracker.pop_next_prompt()
|
|
p3 = self.tracker.pop_next_prompt()
|
|
|
|
self.assertEqual(p1["id"], 1)
|
|
self.assertEqual(p2["id"], 2)
|
|
self.assertEqual(p3["id"], 3)
|
|
|
|
def test_pop_returns_none_when_exhausted(self):
|
|
self.tracker.set_current_prompt("only", {"prompt_id": 1})
|
|
self.tracker.pop_next_prompt()
|
|
result = self.tracker.pop_next_prompt()
|
|
self.assertIsNone(result)
|
|
|
|
def test_queue_preserves_prompt_text(self):
|
|
self.tracker.set_current_prompt("hello world", {"prompt_id": 5})
|
|
ctx = self.tracker.pop_next_prompt()
|
|
self.assertEqual(ctx["text"], "hello world")
|
|
|
|
def test_batch_simulation(self):
|
|
"""Simulate a batch workflow: encode 10 prompts, then pop 10 in order."""
|
|
ids = list(range(100, 110))
|
|
for pid in ids:
|
|
self.tracker.set_current_prompt(f"prompt_{pid}", {"prompt_id": pid})
|
|
|
|
popped_ids = []
|
|
while True:
|
|
ctx = self.tracker.pop_next_prompt()
|
|
if ctx is None:
|
|
break
|
|
popped_ids.append(ctx["id"])
|
|
|
|
self.assertEqual(popped_ids, ids)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|