Vito Sansevero
b1bfce702d
The first push of the branch ran the Tests matrix on all three operating systems. Every Python job failed; the causes were: - numpy missing in CI: the WD14 engine tests import it (requirements-dev) - macOS: temp directories resolve through /private; the admin API tests now use the real path like the gallery roots they compare against - Windows: the startup banner printed an emoji, which raises UnicodeEncodeError on a legacy-code-page console and aborted loading the whole node; the banner is ASCII now and a test loads the package under a cp1252 stdout - Windows: gallery subfolders came back with backslashes, which neither the folder tree nor the search filter understand; they are "/" joined on every platform. A relative DEFAULT_DB_PATH kept its forward slashes after being joined to the extension root; it is normalised now - Windows: sqlite connections opened with `with sqlite3.connect()` are not closed by the context manager and linger until garbage collection, which holds a file lock there. Production sites (vacuum, diagnostics, restore inspection) close explicitly; a restore that swaps the API's database object closes the previous one; PromptModel keeps a registry so a shutdown or a test can close every instance - Windows: served files were still open when their temp directory was removed (tests now read the body), a log download test compared text mode newlines, two tests compared paths with a hard-coded separator tests/open_handles.py makes the Linux run fail the way Windows would: every database test base checks, right before deleting its files, that nothing in the process still holds them open.
283 lines
12 KiB
Python
283 lines
12 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)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|