Vito
d03d0ed1c5
* fix(history): count cached re-runs of runtime-linked text as tracked runs (#403) Since #347 a cached re-run of runtime-linked prompt text (PreviewAny or an LLM into a PromptManager text input) counts in the Library, but its history capture had no raw text, so the history page's "tracked runs" stayed at the number of jobs in which the node executed. The #347 receipt now also keeps the raw, prepend and append text the node gave the history capture. When a cached node's graph inference leaves a field unknown, the capture takes it from the receipt that matches the cache entry, labeled inferred/graph_semantics, so the job joins the prompt's history scope and gets its prompt hint. No schema change and no new evidence vocabulary; no backfill. Fails closed on no receipt, a disagreeing known field, equal outputs with different fields, oversized or invalid text and a full budget (least recently used receipts give up their fields first; counting stays). * docs(testing): link the whitespace history gap to #438
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):
|
|
from unittest import 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()
|