diff --git a/tests/test_prompt_tracker.py b/tests/test_prompt_tracker.py index 2b8c6b2..a9bfdc8 100644 --- a/tests/test_prompt_tracker.py +++ b/tests/test_prompt_tracker.py @@ -34,6 +34,8 @@ def _make_tracker(prompt_timeout=600, cleanup_interval=300): 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 @@ -256,5 +258,59 @@ class TestCleanupExpiredPrompts(unittest.TestCase): 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() diff --git a/tests/test_search_filters.py b/tests/test_search_filters.py new file mode 100644 index 0000000..199b709 --- /dev/null +++ b/tests/test_search_filters.py @@ -0,0 +1,137 @@ +""" +Tests for PromptSearchList output filtering. + +Tests the newline collapsing, multi-part (Clip_) filtering, +and LoRA-only filtering applied to search results. +""" + +import os +import re +import sys +import unittest + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + + +# --- Helpers: replicate the filter logic from prompt_search_list.py --- + +LORA_ONLY = re.compile(r"^(\s*]+>\s*)+$") + + +def apply_filters(texts, skip_multipart=True): + """Apply the same filtering chain as PromptSearchList.search().""" + # Collapse newlines + result = [" ".join(t.split()) for t in texts if t] + + # Filter multi-part + if skip_multipart: + result = [p for p in result if not re.search(r"Clip_\d+", p)] + + # Filter LoRA-only + result = [p for p in result if not LORA_ONLY.match(p)] + + return result + + +class TestNewlineCollapse(unittest.TestCase): + """Newlines must be collapsed so StringOutputList treats each DB entry as one prompt.""" + + def test_single_line_unchanged(self): + result = apply_filters(["a beautiful landscape"]) + self.assertEqual(result, ["a beautiful landscape"]) + + def test_multiline_collapsed_to_spaces(self): + result = apply_filters(["line one\nline two\nline three"]) + self.assertEqual(result, ["line one line two line three"]) + + def test_tabs_and_extra_spaces_collapsed(self): + result = apply_filters(["word1 \t word2\n\nword3"]) + self.assertEqual(result, ["word1 word2 word3"]) + + def test_leading_trailing_whitespace_stripped(self): + result = apply_filters([" padded prompt \n"]) + self.assertEqual(result, ["padded prompt"]) + + +class TestMultipartFilter(unittest.TestCase): + """Prompts containing Clip_N markers should be filtered when skip_multipart=True.""" + + def test_clip_markers_filtered(self): + texts = [ + "Clip_1 a video prompt Clip_2 continuation", + "a normal prompt", + ] + result = apply_filters(texts, skip_multipart=True) + self.assertEqual(result, ["a normal prompt"]) + + def test_clip_markers_kept_when_disabled(self): + texts = ["Clip_1 a video prompt Clip_2 continuation"] + result = apply_filters(texts, skip_multipart=False) + self.assertEqual(len(result), 1) + + def test_clip_in_word_filtered(self): + # "Clip_3" anywhere in text should match + texts = ["some text with Clip_3 marker"] + result = apply_filters(texts, skip_multipart=True) + self.assertEqual(result, []) + + def test_clip_without_number_not_filtered(self): + texts = ["clip art style painting"] + result = apply_filters(texts, skip_multipart=True) + self.assertEqual(result, ["clip art style painting"]) + + +class TestLoraOnlyFilter(unittest.TestCase): + """Prompts that are only LoRA tags with no actual content should be filtered.""" + + def test_single_lora_tag_filtered(self): + result = apply_filters([""]) + self.assertEqual(result, []) + + def test_multiple_lora_tags_filtered(self): + result = apply_filters([" "]) + self.assertEqual(result, []) + + def test_lora_with_content_kept(self): + result = apply_filters([" a beautiful painting"]) + self.assertEqual(result, [" a beautiful painting"]) + + def test_content_with_lora_in_middle_kept(self): + result = apply_filters(["masterpiece, , best quality"]) + self.assertEqual(result, ["masterpiece, , best quality"]) + + def test_lora_with_newlines_collapsed_then_filtered(self): + # After collapsing newlines, this becomes a single line of LoRA tags + result = apply_filters(["\n"]) + self.assertEqual(result, []) + + def test_empty_string_filtered(self): + result = apply_filters([""]) + self.assertEqual(result, []) + + +class TestFilterCombination(unittest.TestCase): + """Test all filters working together on a mixed batch.""" + + def test_mixed_batch(self): + texts = [ + "a good prompt", + "", + "Clip_1\nfirst segment\nClip_2\nsecond segment", + " nice painting, flowers", + "another\ngood\nprompt", + " ", + ] + result = apply_filters(texts) + self.assertEqual( + result, + [ + "a good prompt", + " nice painting, flowers", + "another good prompt", + ], + ) + + +if __name__ == "__main__": + unittest.main()