Vito
d88cf9e7c6
* fix(quality): preserve cursor iteration and option validation Keep prefetched rows on ManagedCursor so partial loops and explicit fetches share one position, and stop existing iterators after close or retirement. Preserve chunked reads and native SQLite result lifetime on reexecution and errors. Catch OptionError before ValueError so legacy auto-tag requests receive the actual validation error. Remove the overwritten ViewerJS title option. Regression tests failed before the fixes for lost rows, reads after close/retirement, and misleading option errors. Focused Python checks pass; cursor checks also pass on Python 3.13 and 3.14. Refs #321. * fix(quality): make the test doubles legal raises and hashes The fake queue no longer raises an attribute that starts as None. The oversized kind records hashes instead of raising from __hash__, which also catches a hash that production code swallows. Hostile.__bool__ raises the TypeError a truth test may raise. * fix(quality): drop the extra saveCoverage argument * refactor(quality): remove assignments that are overwritten before use * refactor(quality): remove the remaining overwritten test assignments * refactor(quality): drop imports repeated inside functions * refactor(quality): give duplicate identity keys one shape * refactor(quality): name the unused job event loop variable * refactor(quality): release model gates in finally, not except BaseException * refactor(quality): end the job runners with an explicit return * refactor(quality): use the precise unittest asserts * refactor(quality): import unittest.mock one way in the tests * fix(quality): say why each empty except is safe, log the broad ones * refactor(quality): stop binding values the tests never read * refactor(quality): remove the unused preview aliases and tag page query * refactor(quality): remove re-exports nothing imports * refactor(quality): parenthesize deliberate string joins in lists * refactor(quality): remove unused JavaScript variables and imports * refactor(quality): use the precise assert in the model status test
333 lines
14 KiB
Python
333 lines
14 KiB
Python
"""Tests for prompt usage tracking: run_count, last_used_at and server-side sorting."""
|
|
|
|
import os
|
|
import sqlite3
|
|
from contextlib import closing
|
|
import sys
|
|
import tempfile
|
|
import unittest
|
|
|
|
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
|
|
try:
|
|
from tests.open_handles import assert_closed # noqa: E402
|
|
except ImportError: # discovered with tests/ as the top-level directory
|
|
from open_handles import assert_closed # noqa: E402
|
|
|
|
from database.models import PromptModel
|
|
from database.operations import PromptDatabase
|
|
from prompt_manager_base import PromptManagerBase
|
|
from utils.hashing import generate_prompt_hash
|
|
|
|
|
|
class UsageTestCase(unittest.TestCase):
|
|
def setUp(self):
|
|
fd, self.path = tempfile.mkstemp(suffix=".db")
|
|
os.close(fd)
|
|
self.db = PromptDatabase(self.path)
|
|
|
|
def tearDown(self):
|
|
if getattr(self, "db", None) is not None:
|
|
self.db.close_all()
|
|
PromptModel.close_all_instances()
|
|
assert_closed(self, self.path)
|
|
for suffix in ("", "-wal", "-shm"):
|
|
if os.path.exists(self.path + suffix):
|
|
os.unlink(self.path + suffix)
|
|
|
|
def _save(self, text, **kwargs):
|
|
return self.db.save_prompt(
|
|
text=text, prompt_hash=generate_prompt_hash(text), **kwargs
|
|
)
|
|
|
|
def _set(self, prompt_id, **columns):
|
|
with closing(sqlite3.connect(self.path)) as conn, conn:
|
|
for column, value in columns.items():
|
|
conn.execute(
|
|
f"UPDATE prompts SET {column} = ? WHERE id = ?", (value, prompt_id)
|
|
)
|
|
|
|
def _ids(self, prompts):
|
|
return [p["id"] for p in prompts]
|
|
|
|
|
|
class TestRecordPromptUse(UsageTestCase):
|
|
def test_new_prompt_starts_unused_and_unstamped(self):
|
|
# Saving is not using: only record_prompt_use stamps last_used_at
|
|
prompt = self.db.get_prompt_by_id(self._save("fresh"))
|
|
self.assertEqual(prompt["run_count"], 0)
|
|
self.assertIsNone(prompt["last_used_at"])
|
|
|
|
def test_never_run_prompts_sort_after_run_ones(self):
|
|
ran = self._save("run once")
|
|
later = self._save("saved later, never run")
|
|
self.db.record_prompt_use(ran)
|
|
last = self._save("saved last, never run")
|
|
|
|
result = self.db.get_recent_prompts(limit=10, sort="last_used_desc")
|
|
# The run prompt leads; never-run prompts follow, newest first
|
|
self.assertEqual(self._ids(result["prompts"]), [ran, last, later])
|
|
|
|
searched = self.db.search_prompts(text="run", sort="last_used_desc")
|
|
self.assertEqual(self._ids(searched)[0], ran)
|
|
|
|
def test_most_used_sort_breaks_ties_with_unstamped_last(self):
|
|
a = self._save("a")
|
|
b = self._save("b")
|
|
self._set(a, run_count=2, last_used_at="2026-01-01T00:00:00.000+00:00")
|
|
self._set(b, run_count=2) # same count, never stamped
|
|
result = self.db.get_recent_prompts(limit=10, sort="run_count_desc")
|
|
self.assertEqual(self._ids(result["prompts"]), [a, b])
|
|
|
|
def test_record_use_increments_and_touches_last_used(self):
|
|
pid = self._save("used")
|
|
self._set(pid, last_used_at="2020-01-01T00:00:00.000+00:00")
|
|
|
|
self.assertTrue(self.db.record_prompt_use(pid))
|
|
self.assertTrue(self.db.record_prompt_use(pid))
|
|
|
|
prompt = self.db.get_prompt_by_id(pid)
|
|
self.assertEqual(prompt["run_count"], 2)
|
|
self.assertGreater(prompt["last_used_at"], "2020-01-01T00:00:00.000+00:00")
|
|
|
|
def test_rerun_right_after_another_save_still_sorts_first(self):
|
|
# No hand-set timestamps: creation and re-run happen within the same millisecond
|
|
old = self._save("old")
|
|
self._save("new")
|
|
self.db.record_prompt_use(old)
|
|
result = self.db.get_recent_prompts(limit=1, sort="last_used_desc")
|
|
self.assertEqual(self._ids(result["prompts"]), [old])
|
|
|
|
def test_microsecond_timestamps_order_correctly_against_backfilled_ones(self):
|
|
backfilled = self._save("backfilled")
|
|
self._set(backfilled, last_used_at="2026-01-01T00:00:00.000+00:00")
|
|
recent = self._save("recent")
|
|
self.db.record_prompt_use(recent)
|
|
self.assertRegex(
|
|
self.db.get_prompt_by_id(recent)["last_used_at"],
|
|
r"^\d{4}-\d\d-\d\dT\d\d:\d\d:\d\d\.\d{6}\+00:00$",
|
|
)
|
|
result = self.db.get_recent_prompts(limit=2, sort="last_used_desc")
|
|
self.assertEqual(self._ids(result["prompts"]), [recent, backfilled])
|
|
|
|
def test_record_use_can_add_several_runs_at_once(self):
|
|
pid = self._save("several")
|
|
self.db.record_prompt_use(pid, times=3)
|
|
self.assertEqual(self.db.get_prompt_by_id(pid)["run_count"], 3)
|
|
|
|
def test_record_use_on_missing_prompt_returns_false(self):
|
|
self.assertFalse(self.db.record_prompt_use(999999))
|
|
|
|
|
|
class TestServerSideSort(UsageTestCase):
|
|
def setUp(self):
|
|
super().setUp()
|
|
self.old = self._save("old prompt", rating=5)
|
|
self.mid = self._save("middle prompt", rating=1)
|
|
self.new = self._save("new prompt")
|
|
self._set(self.old, last_used_at="2026-01-01T00:00:00.000+00:00", run_count=9)
|
|
self._set(self.mid, last_used_at="2026-01-02T00:00:00.000+00:00", run_count=1)
|
|
self._set(self.new, last_used_at="2026-01-03T00:00:00.000+00:00", run_count=3)
|
|
|
|
def test_rerun_bubbles_old_prompt_to_top_across_pages(self):
|
|
self.db.record_prompt_use(self.old)
|
|
page1 = self.db.get_recent_prompts(limit=1, offset=0, sort="last_used_desc")
|
|
self.assertEqual(self._ids(page1["prompts"]), [self.old])
|
|
|
|
def test_most_used(self):
|
|
result = self.db.get_recent_prompts(limit=10, sort="run_count_desc")
|
|
self.assertEqual(self._ids(result["prompts"]), [self.old, self.new, self.mid])
|
|
|
|
def test_rating_sort_puts_unrated_last_and_spans_pages(self):
|
|
page1 = self.db.get_recent_prompts(limit=1, offset=0, sort="rating_desc")
|
|
self.assertEqual(self._ids(page1["prompts"]), [self.old])
|
|
result = self.db.get_recent_prompts(limit=10, sort="rating_asc")
|
|
self.assertEqual(self._ids(result["prompts"]), [self.mid, self.old, self.new])
|
|
|
|
def test_text_sort(self):
|
|
result = self.db.get_recent_prompts(limit=10, sort="text_asc")
|
|
self.assertEqual(self._ids(result["prompts"]), [self.mid, self.new, self.old])
|
|
|
|
def test_default_and_unknown_sort_keep_newest_first(self):
|
|
expected = [self.new, self.mid, self.old]
|
|
self.assertEqual(
|
|
self._ids(self.db.get_recent_prompts(limit=10)["prompts"]), expected
|
|
)
|
|
injected = self.db.get_recent_prompts(limit=10, sort="id; DROP TABLE prompts")
|
|
self.assertEqual(self._ids(injected["prompts"]), expected)
|
|
|
|
def test_search_accepts_sort(self):
|
|
self.db.record_prompt_use(self.mid)
|
|
results = self.db.search_prompts(text="prompt", sort="last_used_desc")
|
|
self.assertEqual(self._ids(results)[0], self.mid)
|
|
|
|
|
|
@unittest.skipIf(
|
|
sqlite3.sqlite_version_info < (3, 35, 0),
|
|
"simulating the old schema needs ALTER TABLE DROP COLUMN (SQLite 3.35+)",
|
|
)
|
|
class TestUsageMigration(unittest.TestCase):
|
|
"""Databases from before 3.2.4 get the columns and an estimated backfill."""
|
|
|
|
def setUp(self):
|
|
fd, self.path = tempfile.mkstemp(suffix=".db")
|
|
os.close(fd)
|
|
self.dbs = []
|
|
db = self._open()
|
|
self.with_images = db.save_prompt(text="has images", prompt_hash="h1")
|
|
self.no_images = db.save_prompt(text="no images", prompt_hash="h2")
|
|
db.close_all()
|
|
with closing(sqlite3.connect(self.path)) as conn, conn:
|
|
conn.execute(
|
|
"UPDATE prompts SET created_at = '2026-01-01T10:00:00.000000+00:00'"
|
|
)
|
|
# generated_images uses SQLite's space-separated CURRENT_TIMESTAMP format
|
|
for i, ts in enumerate(("2026-03-05 08:00:00", "2026-02-01 09:00:00")):
|
|
conn.execute(
|
|
"INSERT INTO generated_images"
|
|
" (prompt_id, image_path, filename, generation_time)"
|
|
" VALUES (?, ?, ?, ?)",
|
|
(self.with_images, f"/out/{i}.png", f"{i}.png", ts),
|
|
)
|
|
# Simulate a pre-3.2.4 schema
|
|
for index in ("idx_prompts_last_used", "idx_prompts_run_count"):
|
|
conn.execute(f"DROP INDEX IF EXISTS {index}")
|
|
conn.execute("ALTER TABLE prompts DROP COLUMN last_used_at")
|
|
conn.execute("ALTER TABLE prompts DROP COLUMN run_count")
|
|
# Schema init runs once per process per path; the next PromptDatabase()
|
|
# must look at the file again to see the downgraded schema.
|
|
PromptModel.reset_schema_cache()
|
|
|
|
def _open(self):
|
|
"""Open the database under test and close it again in tearDown."""
|
|
db = PromptDatabase(self.path)
|
|
self.dbs.append(db)
|
|
return db
|
|
|
|
def tearDown(self):
|
|
for db in self.dbs:
|
|
db.close_all()
|
|
PromptModel.close_all_instances()
|
|
assert_closed(self, self.path)
|
|
for suffix in ("", "-wal", "-shm"):
|
|
if os.path.exists(self.path + suffix):
|
|
os.unlink(self.path + suffix)
|
|
|
|
def test_backfill_estimates_usage_from_linked_images(self):
|
|
db = self._open()
|
|
used = db.get_prompt_by_id(self.with_images)
|
|
unused = db.get_prompt_by_id(self.no_images)
|
|
|
|
self.assertEqual(used["run_count"], 2)
|
|
self.assertTrue(used["last_used_at"].startswith("2026-03-05T08:00:00"))
|
|
self.assertEqual(unused["run_count"], 1)
|
|
self.assertTrue(unused["last_used_at"].startswith("2026-01-01T10:00:00"))
|
|
|
|
def test_backfill_runs_only_when_the_columns_are_added(self):
|
|
db = self._open()
|
|
# A row whose usage is unknown (e.g. written by 3.2.3 after a downgrade)
|
|
# stays "never run": a restart must not invent a run for it, otherwise
|
|
# Recently Used could never tell saved-but-unused prompts from run ones.
|
|
with closing(sqlite3.connect(self.path)) as conn, conn:
|
|
conn.execute(
|
|
"INSERT INTO prompts (text, hash, created_at)"
|
|
" VALUES ('from 3.2.3', 'h3', '2026-04-01T00:00:00.000000+00:00')"
|
|
)
|
|
legacy_id = conn.execute(
|
|
"SELECT id FROM prompts WHERE hash = 'h3'"
|
|
).fetchone()[0]
|
|
self.assertIsNone(db.get_prompt_by_id(legacy_id)["last_used_at"])
|
|
|
|
PromptModel.reset_schema_cache()
|
|
later = self._open().get_prompt_by_id(legacy_id)
|
|
self.assertIsNone(later["last_used_at"])
|
|
self.assertEqual(later["run_count"], 0)
|
|
|
|
def test_migration_runs_once(self):
|
|
db = self._open()
|
|
db.record_prompt_use(self.no_images)
|
|
reopened = self._open()
|
|
self.assertEqual(reopened.get_prompt_by_id(self.no_images)["run_count"], 2)
|
|
|
|
|
|
class TestNodeRecordsUse(UsageTestCase):
|
|
"""Node saves count a run only when asked; re-runs are counted by the queue hook.
|
|
|
|
ComfyUI skips unchanged nodes, so counting at execution misses re-runs; see
|
|
tests/test_usage_tracking.py for the hook that counts them.
|
|
"""
|
|
|
|
def _node(self):
|
|
node = PromptManagerBase.__new__(PromptManagerBase)
|
|
node.db = self.db
|
|
node.logger = __import__("logging").getLogger("test.prompt_usage")
|
|
return node
|
|
|
|
def test_save_without_count_run_leaves_counting_to_the_queue_hook(self):
|
|
node = self._node()
|
|
first = node._save_prompt_to_database("a lighthouse at dusk")
|
|
second = node._save_prompt_to_database("a lighthouse at dusk")
|
|
|
|
self.assertEqual(first, second)
|
|
self.assertEqual(self.db.get_prompt_by_id(first)["run_count"], 0)
|
|
|
|
def test_count_run_counts_new_and_existing_prompts(self):
|
|
node = self._node()
|
|
pid = node._save_prompt_to_database("batch item", count_run=True)
|
|
node._save_prompt_to_database("batch item", count_run=True)
|
|
self.assertEqual(self.db.get_prompt_by_id(pid)["run_count"], 2)
|
|
|
|
|
|
class TestNodeRecordsCachedRunReceipt(UsageTestCase):
|
|
"""Each execution leaves a receipt for cached jobs (#347)."""
|
|
|
|
GRAPH = {
|
|
"1": {"class_type": "PromptManagerText", "inputs": {"text": "typed text"}},
|
|
}
|
|
|
|
def _node(self):
|
|
import unittest.mock as mock
|
|
|
|
from prompt_manager_text import PromptManagerText
|
|
|
|
node = PromptManagerText.__new__(PromptManagerText)
|
|
node.db = self.db
|
|
node.logger = __import__("logging").getLogger("test.prompt_usage")
|
|
node.prompt_tracker = mock.Mock()
|
|
node.comfyui_integration = mock.Mock()
|
|
node._inject_lora_trigger_words = lambda text: text
|
|
node.cached_runs = mock.Mock()
|
|
node.cached_runs.replayed.return_value = False
|
|
return node
|
|
|
|
def test_receipt_holds_the_returned_text_and_the_stored_hash(self):
|
|
from utils.cached_runs import Facts
|
|
|
|
node = self._node()
|
|
node.process_text(
|
|
"typed text", prepend_text="pre", prompt=self.GRAPH, unique_id="1"
|
|
)
|
|
node.cached_runs.record_item.assert_called_once_with(
|
|
"1",
|
|
("pre typed text",),
|
|
Facts(generate_prompt_hash("typed text")),
|
|
# #403: the texts the history capture got, for a cached job's history
|
|
(("typed text", "pre", ""),),
|
|
)
|
|
# Hook-visible text: the node counts nothing itself
|
|
row = self.db.get_prompt_by_hash(generate_prompt_hash("typed text"))
|
|
self.assertEqual(row["run_count"], 0)
|
|
|
|
def test_empty_editor_records_empty_facts(self):
|
|
from utils.cached_runs import Facts
|
|
|
|
node = self._node()
|
|
node.process_text(" ", prompt=self.GRAPH, unique_id="1")
|
|
node.cached_runs.record_item.assert_called_once_with(
|
|
"1", (" ",), Facts(None), ((" ", "", ""),)
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|