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
1505 lines
63 KiB
Python
1505 lines
63 KiB
Python
"""#342: recorded history attributes a file only to one certain, stored prompt."""
|
|
|
|
import json
|
|
import os
|
|
import sqlite3
|
|
import threading
|
|
import time
|
|
import unittest
|
|
from pathlib import Path
|
|
import unittest.mock as mock
|
|
|
|
from database.run_attribution import (
|
|
RecordedOutput,
|
|
recorded_output,
|
|
stored_run_prompt,
|
|
)
|
|
from tests import test_execution_history_database as history_tests
|
|
from tests.background_stops import close_handler
|
|
from tests.test_api_admin import AdminAPITestCase
|
|
from tests import saver_anchor_graphs as shapes
|
|
from tests.test_pair_scan import sampled_graph
|
|
from tests.test_execution_context_database import known, snapshot
|
|
from tests.test_prompt_attribution import issue_graph, llm_text_graph, make_tracker
|
|
from utils.execution_contexts import (
|
|
ContextBundle,
|
|
PromptSnapshot,
|
|
RecipeContext,
|
|
TextEvidence,
|
|
)
|
|
from utils.execution_outputs import (
|
|
OutputBundle,
|
|
OutputReference,
|
|
snapshot_output_roots,
|
|
)
|
|
from utils.hashing import generate_prompt_hash
|
|
from py.config import GalleryConfig
|
|
from utils.image_monitor import ImageGenerationHandler
|
|
|
|
STYLE = "windows" if os.name == "nt" else "posix"
|
|
LLM_TEXT = "Photorealistic aerial landscape of misty mountains at dawn"
|
|
|
|
|
|
def now_ms():
|
|
return time.time_ns() // 1_000_000
|
|
|
|
|
|
class RecordedRuns:
|
|
"""Write runs into execution history the way the capture pipeline does."""
|
|
|
|
start = history_tests.HistoryDatabaseTests.start
|
|
finish = history_tests.HistoryDatabaseTests.finish
|
|
|
|
def record_run(
|
|
self, db, output_root, files, positives, item=0, window=None, nodes=None
|
|
):
|
|
started, finished = window or (now_ms() - 5000, now_ms() + 5000)
|
|
start = self.start(db, item, started)
|
|
db.write_execution_start(start)
|
|
finish = self.finish(start, time=finished)
|
|
db.write_execution_finish(finish)
|
|
if positives: # a run with no PromptManager node records no contexts
|
|
db.write_execution_contexts(
|
|
start.run_id,
|
|
start.storage_epoch,
|
|
ContextBundle(
|
|
tuple(
|
|
RecipeContext(
|
|
str(4 + i), snapshot(text), snapshot(""), "inferred"
|
|
)
|
|
for i, text in enumerate(positives)
|
|
),
|
|
"complete",
|
|
"captured",
|
|
),
|
|
)
|
|
root = snapshot_output_roots(STYLE, output_root, output_root).output_key
|
|
nodes = nodes or ["9"] * len(files)
|
|
references = tuple(
|
|
OutputReference(node, i, "output", sub, name, STYLE, root)
|
|
for i, (node, (sub, name)) in enumerate(zip(nodes, files))
|
|
)
|
|
db.write_execution_outputs_for_finish(
|
|
finish, OutputBundle(references, "complete", "reported")
|
|
)
|
|
return start.run_id
|
|
|
|
|
|
class TestStoredRunPrompt(RecordedRuns, unittest.TestCase):
|
|
close_databases = history_tests.HistoryDatabaseTests.close_databases
|
|
database = history_tests.HistoryDatabaseTests.database
|
|
|
|
def setUp(self): # the shared temp dir and database list first
|
|
history_tests.HistoryDatabaseTests.setUp(self)
|
|
self.db = self.database()
|
|
self.output = str(Path(self.temp.name) / "output")
|
|
os.makedirs(os.path.join(self.output, "sub"))
|
|
self.image = os.path.join(self.output, "sub", "img_00001_.png")
|
|
Path(self.image).write_bytes(b"png")
|
|
|
|
def save(self, text):
|
|
return self.db.save_prompt(text=text, prompt_hash=generate_prompt_hash(text))
|
|
|
|
def attribute(self, path=None):
|
|
return stored_run_prompt(self.db, path or self.image, self.output)
|
|
|
|
def test_one_run_naming_one_stored_prompt_attributes_the_file(self):
|
|
prompt_id = self.save("an llm prompt")
|
|
self.record_run(
|
|
self.db, self.output, [("sub", "img_00001_.png")], ["an llm prompt"]
|
|
)
|
|
self.assertEqual(self.attribute(), prompt_id)
|
|
|
|
def test_text_that_is_not_stored_attributes_nothing(self):
|
|
"""The #342 run: raw editor text was empty, nothing was stored."""
|
|
self.save("a woman holding a coffee cup")
|
|
self.record_run(self.db, self.output, [("sub", "img_00001_.png")], [""])
|
|
self.assertIsNone(self.attribute())
|
|
|
|
def test_two_runs_reporting_the_same_file_attribute_nothing(self):
|
|
self.save("first")
|
|
self.save("second")
|
|
self.record_run(self.db, self.output, [("sub", "img_00001_.png")], ["first"])
|
|
self.record_run(
|
|
self.db, self.output, [("sub", "img_00001_.png")], ["second"], item=1
|
|
)
|
|
self.assertIsNone(self.attribute())
|
|
|
|
def test_a_file_written_outside_the_run_attributes_nothing(self):
|
|
self.save("an llm prompt")
|
|
self.record_run(
|
|
self.db,
|
|
self.output,
|
|
[("sub", "img_00001_.png")],
|
|
["an llm prompt"],
|
|
window=(1_000_000, 2_000_000),
|
|
)
|
|
self.assertIsNone(self.attribute())
|
|
|
|
def test_contexts_naming_different_prompts_attribute_nothing(self):
|
|
self.save("left")
|
|
self.save("right")
|
|
self.record_run(
|
|
self.db, self.output, [("sub", "img_00001_.png")], ["left", "right"]
|
|
)
|
|
self.assertIsNone(self.attribute())
|
|
|
|
def test_files_outside_the_output_root_or_unrecorded_attribute_nothing(self):
|
|
self.save("an llm prompt")
|
|
self.record_run(
|
|
self.db, self.output, [("sub", "img_00001_.png")], ["an llm prompt"]
|
|
)
|
|
outside = os.path.join(self.temp.name, "img_00001_.png")
|
|
Path(outside).write_bytes(b"png")
|
|
self.assertIsNone(self.attribute(outside))
|
|
other = os.path.join(self.output, "sub", "other.png")
|
|
Path(other).write_bytes(b"png")
|
|
self.assertIsNone(self.attribute(other))
|
|
self.assertIsNone(stored_run_prompt(self.db, self.image, None))
|
|
self.assertIsNone(stored_run_prompt(self.db, self.image, "relative/root"))
|
|
|
|
def test_a_library_without_history_tables_attributes_nothing(self):
|
|
broken = mock.Mock()
|
|
broken.model.get_connection.return_value.execute.side_effect = (
|
|
sqlite3.OperationalError("no such table")
|
|
)
|
|
self.assertIsNone(stored_run_prompt(broken, self.image, self.output))
|
|
|
|
def test_the_monitor_links_from_history_when_nothing_else_is_certain(self):
|
|
prompt_id = self.save(LLM_TEXT)
|
|
self.record_run(self.db, self.output, [("sub", "img_00001_.png")], [LLM_TEXT])
|
|
handler = ImageGenerationHandler(self.db, make_tracker(self.db))
|
|
self.addCleanup(close_handler, self, handler)
|
|
linked = []
|
|
handler.link_image_to_prompt = lambda path, prompt, meta: linked.append(
|
|
prompt["id"]
|
|
)
|
|
with (
|
|
mock.patch(
|
|
"utils.image_monitor.host_output_directory", return_value=self.output
|
|
),
|
|
mock.patch.object(
|
|
handler.metadata_extractor,
|
|
"extract_metadata",
|
|
return_value={"prompt": llm_text_graph()},
|
|
),
|
|
):
|
|
handler.process_new_image(self.image)
|
|
self.assertEqual(linked, [prompt_id])
|
|
|
|
|
|
class TestRecordedOutputNode(RecordedRuns, unittest.TestCase):
|
|
"""#386: the save node history says wrote a file, under the same rules."""
|
|
|
|
close_databases = history_tests.HistoryDatabaseTests.close_databases
|
|
database = history_tests.HistoryDatabaseTests.database
|
|
setUp = TestStoredRunPrompt.setUp
|
|
|
|
def node(self, path=None):
|
|
found = recorded_output(self.db, path or self.image, self.output)
|
|
return found.saver_id if found is not None else None
|
|
|
|
def test_one_run_one_saver_names_the_saver(self):
|
|
other = os.path.join(self.output, "sub", "img_00002_.png")
|
|
Path(other).write_bytes(b"png")
|
|
self.record_run(
|
|
self.db,
|
|
self.output,
|
|
[("sub", "img_00001_.png"), ("sub", "img_00002_.png")],
|
|
[],
|
|
nodes=["13", "23"],
|
|
)
|
|
self.assertEqual(self.node(), "13")
|
|
self.assertEqual(self.node(other), "23")
|
|
|
|
def test_a_run_without_prompt_manager_contexts_still_names_the_saver(self):
|
|
self.record_run(self.db, self.output, [("sub", "img_00001_.png")], [])
|
|
self.assertEqual(self.node(), "9")
|
|
|
|
def test_two_runs_reporting_the_file_name_nothing(self):
|
|
self.record_run(
|
|
self.db, self.output, [("sub", "img_00001_.png")], [], nodes=["13"]
|
|
)
|
|
self.record_run(
|
|
self.db, self.output, [("sub", "img_00001_.png")], [], item=1, nodes=["23"]
|
|
)
|
|
self.assertIsNone(self.node())
|
|
|
|
def test_one_run_reporting_the_file_from_two_savers_names_nothing(self):
|
|
self.record_run(
|
|
self.db,
|
|
self.output,
|
|
[("sub", "img_00001_.png"), ("sub", "img_00001_.png")],
|
|
[],
|
|
nodes=["13", "23"],
|
|
)
|
|
self.assertIsNone(self.node())
|
|
|
|
def test_a_file_written_outside_the_run_names_nothing(self):
|
|
self.record_run(
|
|
self.db,
|
|
self.output,
|
|
[("sub", "img_00001_.png")],
|
|
[],
|
|
window=(1_000_000, 2_000_000),
|
|
)
|
|
self.assertIsNone(self.node())
|
|
|
|
def test_unrecorded_outside_or_unknown_roots_name_nothing(self):
|
|
self.record_run(self.db, self.output, [("sub", "img_00001_.png")], [])
|
|
other = os.path.join(self.output, "sub", "other.png")
|
|
Path(other).write_bytes(b"png")
|
|
self.assertIsNone(self.node(other))
|
|
self.assertIsNone(recorded_output(self.db, self.image, None))
|
|
self.assertIsNone(recorded_output(self.db, "", self.output))
|
|
|
|
def test_the_run_texts_recorded_at_execution_come_with_the_saver(self):
|
|
self.record_run(
|
|
self.db, self.output, [("sub", "img_00001_.png")], ["an llm prompt"]
|
|
)
|
|
self.assertEqual(
|
|
recorded_output(self.db, self.image, self.output),
|
|
RecordedOutput("9", {"4": {"positive": "an llm prompt", "negative": ""}}),
|
|
)
|
|
|
|
def test_texts_only_inferred_from_the_graph_are_not_recorded(self):
|
|
start = self.start(self.db, 0, now_ms() - 5000)
|
|
self.db.write_execution_start(start)
|
|
finish = self.finish(start, time=now_ms() + 5000)
|
|
self.db.write_execution_finish(finish)
|
|
inferred = TextEvidence("graph text", "inferred", "graph_literal", None)
|
|
self.db.write_execution_contexts(
|
|
start.run_id,
|
|
start.storage_epoch,
|
|
ContextBundle(
|
|
(
|
|
RecipeContext(
|
|
"7",
|
|
PromptSnapshot(inferred, inferred, known(""), known("")),
|
|
snapshot("kept"),
|
|
"inferred",
|
|
),
|
|
),
|
|
"complete",
|
|
"captured",
|
|
),
|
|
)
|
|
root = snapshot_output_roots(STYLE, self.output, self.output).output_key
|
|
self.db.write_execution_outputs_for_finish(
|
|
finish,
|
|
OutputBundle(
|
|
(
|
|
OutputReference(
|
|
"9", 0, "output", "sub", "img_00001_.png", STYLE, root
|
|
),
|
|
),
|
|
"complete",
|
|
"reported",
|
|
),
|
|
)
|
|
self.assertEqual(
|
|
recorded_output(self.db, self.image, self.output).texts,
|
|
{"7": {"positive": None, "negative": "kept"}},
|
|
)
|
|
|
|
def test_several_savers_name_no_saver_but_keep_the_texts(self):
|
|
self.record_run(
|
|
self.db,
|
|
self.output,
|
|
[("sub", "img_00001_.png"), ("sub", "img_00001_.png")],
|
|
["text"],
|
|
nodes=["13", "23"],
|
|
)
|
|
found = recorded_output(self.db, self.image, self.output)
|
|
self.assertIsNone(found.saver_id)
|
|
self.assertEqual(found.texts, {"4": {"positive": "text", "negative": ""}})
|
|
|
|
def test_the_route_lookup_logs_a_failure_and_names_nothing(self):
|
|
from database.run_attribution import recorded_output_lookup
|
|
|
|
broken = mock.Mock()
|
|
broken.model.get_connection.return_value.execute.side_effect = (
|
|
sqlite3.OperationalError("locked")
|
|
)
|
|
broken.get_execution_contexts.side_effect = sqlite3.OperationalError("x")
|
|
with mock.patch(
|
|
"database.run_attribution.host_output_directory", return_value=self.output
|
|
):
|
|
self.assertIsNone(recorded_output_lookup(broken)(self.image))
|
|
with (
|
|
mock.patch(
|
|
"database.run_attribution.recorded_runs",
|
|
side_effect=sqlite3.OperationalError("locked"),
|
|
),
|
|
self.assertLogs("prompt_manager.run_attribution", "WARNING"),
|
|
):
|
|
self.assertIsNone(
|
|
recorded_output_lookup(broken, close=True)(self.image)
|
|
)
|
|
broken.model.close.assert_called_once_with()
|
|
|
|
def test_a_history_reference_spells_the_file_as_history_recorded_it(self):
|
|
# On Windows the history route's target is normcased (lower case);
|
|
# history matches names exactly, so the route passes its reference.
|
|
from database.run_attribution import recorded_output_lookup
|
|
|
|
upper = os.path.join(self.output, "sub", "IMG_00001_.png")
|
|
Path(upper).write_bytes(b"png")
|
|
self.record_run(self.db, self.output, [("sub", "IMG_00001_.png")], [])
|
|
target = os.path.join(self.output, "sub", "img_00001_.png")
|
|
reference = {
|
|
"folder_type": "output",
|
|
"subfolder": "sub",
|
|
"filename": "IMG_00001_.png",
|
|
}
|
|
with mock.patch(
|
|
"database.run_attribution.host_output_directory", return_value=self.output
|
|
):
|
|
found = recorded_output_lookup(self.db, reference=reference)(target)
|
|
self.assertEqual(found.saver_id, "9")
|
|
temp = {**reference, "folder_type": "temp"}
|
|
self.assertIsNone(recorded_output_lookup(self.db, reference=temp)(target))
|
|
|
|
def test_a_library_without_history_tables_names_nothing(self):
|
|
broken = mock.Mock()
|
|
broken.model.get_connection.return_value.execute.side_effect = (
|
|
sqlite3.OperationalError("no such table")
|
|
)
|
|
self.assertIsNone(recorded_output(broken, self.image, self.output))
|
|
|
|
|
|
class TestScanUsesHistory(RecordedRuns, AdminAPITestCase):
|
|
"""The output scan applies the same rules as the monitor."""
|
|
|
|
def linked_prompt(self, path):
|
|
row = (
|
|
self.api.db.model.get_connection()
|
|
.execute(
|
|
"SELECT prompt_id FROM generated_images WHERE image_path=?",
|
|
(str(path),),
|
|
)
|
|
.fetchone()
|
|
)
|
|
return row[0] if row else None
|
|
|
|
def write(self, name):
|
|
path = self.output_dir / name
|
|
path.write_bytes(b"png")
|
|
return path
|
|
|
|
def test_an_llm_text_file_links_to_the_prompt_its_run_stored(self):
|
|
db = self.api.db
|
|
prompt_id = db.save_prompt(
|
|
text=LLM_TEXT, prompt_hash=generate_prompt_hash(LLM_TEXT)
|
|
)
|
|
path = self.write("llm_00001_.png")
|
|
self.record_run(db, str(self.output_dir), [("", path.name)], [LLM_TEXT])
|
|
meta = {"prompt": json.dumps(llm_text_graph())}
|
|
self.assertEqual(self.api._ingest_scanned_file(path, meta), "linked")
|
|
self.assertEqual(self.linked_prompt(path), prompt_id)
|
|
|
|
def test_a_metadata_less_file_links_through_its_recorded_run(self):
|
|
"""I1: --disable-metadata outputs carry no prompt; history still proves one."""
|
|
db = self.api.db
|
|
prompt_id = db.save_prompt(
|
|
text=LLM_TEXT, prompt_hash=generate_prompt_hash(LLM_TEXT)
|
|
)
|
|
path = self.write("bare_00001_.png")
|
|
self.record_run(db, str(self.output_dir), [("", path.name)], [LLM_TEXT])
|
|
self.assertEqual(self.api._ingest_scanned_file(path, {}), "linked")
|
|
self.assertEqual(self.linked_prompt(path), prompt_id)
|
|
self.assertEqual(self.api._ingest_scanned_file(path, None), "unchanged")
|
|
detailed = self.api._ingest_scanned_file(path, {}, detailed=True)
|
|
self.assertEqual(detailed["status"], "unchanged")
|
|
self.assertTrue(detailed["title"])
|
|
|
|
def test_a_metadata_less_file_without_a_recorded_run_is_still_skipped(self):
|
|
path = self.write("bare_00002_.png")
|
|
self.assertIsNone(self.api._ingest_scanned_file(path, {}))
|
|
self.assertIsNone(self.api._ingest_scanned_file(path, {"workflow": {}}))
|
|
self.assertIsNone(self.linked_prompt(path))
|
|
|
|
def test_the_issue_file_stays_unlinked_even_with_its_pair_record(self):
|
|
"""Resolved with an empty stored positive: before #342 the scan raised."""
|
|
db = self.api.db
|
|
db.save_prompt(
|
|
text="a woman holding a coffee cup",
|
|
prompt_hash=generate_prompt_hash("a woman holding a coffee cup"),
|
|
)
|
|
path = self.write("moss_00009_.png")
|
|
self.record_run(db, str(self.output_dir), [("", path.name)], [""])
|
|
evidence = {
|
|
"version": 1,
|
|
"nodes": {
|
|
"18:35": {
|
|
"records": [
|
|
{
|
|
"text": "",
|
|
"negative_text": "",
|
|
"positive_text": LLM_TEXT,
|
|
"negative_effective_text": "",
|
|
"prepend_text": LLM_TEXT,
|
|
"append_text": "",
|
|
}
|
|
]
|
|
}
|
|
},
|
|
}
|
|
meta = {
|
|
"prompt": json.dumps(issue_graph()),
|
|
"comfyassets_prompt_pairs": json.dumps(evidence),
|
|
}
|
|
self.assertEqual(self.api._ingest_scanned_file(path, meta), "unresolved")
|
|
self.assertIsNone(self.linked_prompt(path))
|
|
self.assertEqual(db.list_prompt_pairs()["pairs"], [])
|
|
|
|
def scan_dataset(self, names, nodes):
|
|
"""Two branches of the shared-prefix dataset graph, one run, then a scan."""
|
|
paths = [self.write(name) for name in names]
|
|
self.record_run(
|
|
self.api.db,
|
|
str(self.output_dir),
|
|
[("", path.name) for path in paths],
|
|
[],
|
|
nodes=nodes,
|
|
)
|
|
meta = {"prompt": json.dumps(shapes.shared_prefix())}
|
|
return paths, self.api._ingest_scan_batch_sync([(p, meta) for p in paths])
|
|
|
|
def test_the_recorded_saver_links_each_dataset_file_to_its_branch(self):
|
|
"""#247: before, the legacy parser gave every branch the first prompt."""
|
|
paths, counts = self.scan_dataset(
|
|
["IMG_GEN_00001_.png", "IMG_GEN_00002_.png"], ["13", "23"]
|
|
)
|
|
self.assertEqual(counts["added"], 2, counts)
|
|
for path, text in zip(paths, shapes.BRANCH_TEXTS):
|
|
prompt = self.api.db.get_prompt_by_hash(generate_prompt_hash(text))
|
|
self.assertEqual(self.linked_prompt(path), prompt["id"])
|
|
self.assertEqual(prompt["category"], "scanned")
|
|
|
|
def test_a_file_the_legacy_answer_linked_is_left_as_it_is(self):
|
|
"""#247 Q5: a rescan never adds a second link to the branch prompt."""
|
|
db = self.api.db
|
|
legacy = shapes.BRANCH_TEXTS[0]
|
|
legacy_id = db.save_prompt(legacy, prompt_hash=generate_prompt_hash(legacy))
|
|
path = self.write("IMG_GEN_00002_.png")
|
|
db.link_image_to_prompt(legacy_id, str(path))
|
|
self.record_run(db, str(self.output_dir), [("", path.name)], [], nodes=["23"])
|
|
meta = {"prompt": json.dumps(shapes.shared_prefix())}
|
|
self.assertEqual(self.api._ingest_scanned_file(path, meta), "conflict")
|
|
self.assertEqual(self.linked_prompt(path), legacy_id)
|
|
self.assertIsNone(
|
|
db.get_prompt_by_hash(generate_prompt_hash(shapes.BRANCH_TEXTS[1]))
|
|
)
|
|
|
|
def test_a_recorded_run_never_adds_a_second_link(self):
|
|
db = self.api.db
|
|
db.save_prompt(text=LLM_TEXT, prompt_hash=generate_prompt_hash(LLM_TEXT))
|
|
other = db.save_prompt(text="a dog", prompt_hash=generate_prompt_hash("a dog"))
|
|
path = self.write("bare_00003_.png")
|
|
db.link_image_to_prompt(other, str(path))
|
|
self.record_run(db, str(self.output_dir), [("", path.name)], [LLM_TEXT])
|
|
self.assertEqual(self.api._ingest_scanned_file(path, {}), "conflict")
|
|
rows = db.model.get_connection().execute(
|
|
"SELECT prompt_id FROM generated_images WHERE image_path=?", (str(path),)
|
|
)
|
|
self.assertEqual([row[0] for row in rows], [other])
|
|
|
|
def test_a_single_file_asks_history_once(self):
|
|
db = self.api.db
|
|
db.save_prompt(text=LLM_TEXT, prompt_hash=generate_prompt_hash(LLM_TEXT))
|
|
path = self.write("llm_00002_.png")
|
|
self.record_run(db, str(self.output_dir), [("", path.name)], [LLM_TEXT])
|
|
meta = {"prompt": json.dumps(llm_text_graph())}
|
|
statements = []
|
|
conn = db.model.get_connection()
|
|
conn.set_trace_callback(statements.append)
|
|
try:
|
|
self.assertEqual(self.api._ingest_scanned_file(path, meta), "linked")
|
|
finally:
|
|
conn.set_trace_callback(None)
|
|
lookups = [s for s in statements if "FROM execution_run_outputs" in s]
|
|
self.assertEqual(len(lookups), 1)
|
|
|
|
def test_a_file_the_recorded_saver_links_is_not_queued(self):
|
|
"""A default-ambiguous file linked at the inspector's step is no
|
|
Needs review case."""
|
|
with mock.patch.object(
|
|
type(self.api), "_record_scanned_review", autospec=True
|
|
) as record:
|
|
_paths, counts = self.scan_dataset(
|
|
["IMG_GEN_00003_.png", "IMG_GEN_00004_.png"], ["13", "23"]
|
|
)
|
|
self.assertEqual(counts["added"], 2, counts)
|
|
record.assert_not_called()
|
|
|
|
|
|
class TestHostHelpers(unittest.TestCase):
|
|
"""The host lookups degrade to "no proof", never to an error."""
|
|
|
|
def test_host_output_directory(self):
|
|
import sys
|
|
import types
|
|
|
|
from database import run_attribution
|
|
|
|
cases = (
|
|
(types.SimpleNamespace(get_output_directory=lambda: "/out"), "/out"),
|
|
(types.SimpleNamespace(get_output_directory=lambda: ""), None),
|
|
(types.SimpleNamespace(get_output_directory=lambda: 1 / 0), None),
|
|
)
|
|
original = sys.modules.get("folder_paths")
|
|
try:
|
|
for module, expected in cases:
|
|
sys.modules["folder_paths"] = module
|
|
self.assertEqual(run_attribution.host_output_directory(), expected)
|
|
sys.modules["folder_paths"] = None # import fails: outside ComfyUI
|
|
self.assertIsNone(run_attribution.host_output_directory())
|
|
finally:
|
|
if original is None:
|
|
sys.modules.pop("folder_paths", None)
|
|
else:
|
|
sys.modules["folder_paths"] = original
|
|
|
|
def test_run_window_needs_integer_times_and_a_readable_file(self):
|
|
from database.run_attribution import _written_during
|
|
|
|
self.assertFalse(_written_during(__file__, None, 5))
|
|
self.assertFalse(_written_during(os.path.join("no", "such.png"), 0, 5))
|
|
|
|
def test_executing_run_id(self):
|
|
import sys
|
|
import types
|
|
|
|
from utils.run_keys import executing_run_id
|
|
|
|
def host(context):
|
|
def get_executing_context():
|
|
if isinstance(context, Exception):
|
|
raise context
|
|
return context
|
|
|
|
return types.SimpleNamespace(get_executing_context=get_executing_context)
|
|
|
|
original = sys.modules.get("comfy_execution.utils")
|
|
try:
|
|
for context, expected in (
|
|
(types.SimpleNamespace(prompt_id="run-1"), "run-1"),
|
|
(types.SimpleNamespace(prompt_id=""), None),
|
|
(None, None),
|
|
(RuntimeError("no context"), None),
|
|
):
|
|
sys.modules["comfy_execution.utils"] = host(context)
|
|
self.assertEqual(executing_run_id(), expected)
|
|
sys.modules["comfy_execution.utils"] = None
|
|
self.assertIsNone(executing_run_id())
|
|
finally:
|
|
if original is None:
|
|
sys.modules.pop("comfy_execution.utils", None)
|
|
else:
|
|
sys.modules["comfy_execution.utils"] = original
|
|
|
|
|
|
class TestMonitorHistoryRecheck(RecordedRuns, unittest.TestCase):
|
|
"""I1: a keyless image gets one history-only re-check when its run is sealed."""
|
|
|
|
def setUp(self):
|
|
history_tests.HistoryDatabaseTests.setUp(self)
|
|
self.db = history_tests.HistoryDatabaseTests.database(self)
|
|
self.output = str(Path(self.temp.name) / "output")
|
|
os.makedirs(self.output)
|
|
self.handler = ImageGenerationHandler(self.db, make_tracker(self.db))
|
|
self.handler.processing_delay = 0
|
|
self.handler.settle_interval = 0.01
|
|
self.addCleanup(close_handler, self, self.handler)
|
|
self.linked = []
|
|
self.handler.link_image_to_prompt = lambda path, prompt, meta: (
|
|
self.linked.append((os.path.basename(path), prompt["id"]))
|
|
)
|
|
patches = (
|
|
mock.patch(
|
|
"utils.image_monitor.host_output_directory", return_value=self.output
|
|
),
|
|
mock.patch.object(
|
|
self.handler.metadata_extractor, "extract_metadata", return_value=None
|
|
),
|
|
)
|
|
for patch in patches:
|
|
patch.start()
|
|
self.addCleanup(patch.stop)
|
|
|
|
close_databases = history_tests.HistoryDatabaseTests.close_databases
|
|
|
|
def image(self, name):
|
|
path = os.path.join(self.output, name)
|
|
Path(path).write_bytes(b"png")
|
|
return path
|
|
|
|
def save(self, text):
|
|
return self.db.save_prompt(text=text, prompt_hash=generate_prompt_hash(text))
|
|
|
|
def wait_for(self, predicate, timeout=5.0):
|
|
deadline = time.monotonic() + timeout
|
|
while time.monotonic() < deadline:
|
|
if predicate():
|
|
return True
|
|
time.sleep(0.01)
|
|
return predicate()
|
|
|
|
def test_the_run_ending_after_processing_links_the_image_once(self):
|
|
prompt_id = self.save(LLM_TEXT)
|
|
path = self.image("late_00001_.png")
|
|
self.handler.process_new_image(path) # the run is still executing
|
|
self.assertEqual(self.linked, [])
|
|
self.assertIn(
|
|
os.path.normcase(os.path.realpath(path)), self.handler.awaiting_history()
|
|
)
|
|
self.record_run(self.db, self.output, [("", "late_00001_.png")], [LLM_TEXT])
|
|
self.assertTrue(self.wait_for(lambda: self.linked))
|
|
self.assertTrue(self.wait_for(lambda: self.handler.pending_count() == 0))
|
|
self.assertEqual(self.linked, [("late_00001_.png", prompt_id)])
|
|
self.assertEqual(self.handler.awaiting_history(), [])
|
|
# Another run reporting the same name later is not a second re-check
|
|
self.record_run(
|
|
self.db, self.output, [("", "late_00001_.png")], [LLM_TEXT], item=1
|
|
)
|
|
self.assertTrue(self.wait_for(lambda: self.handler.pending_count() == 0))
|
|
self.assertEqual(len(self.linked), 1)
|
|
|
|
def test_a_seal_during_the_first_pass_is_not_lost(self):
|
|
"""R1: the run seals while the image's own pass is still pending."""
|
|
prompt_id = self.save(LLM_TEXT)
|
|
path = self.image("race_00001_.png")
|
|
real = self.handler._history_prompt
|
|
sealed = []
|
|
|
|
def history(image_path):
|
|
result = real(image_path) # before the run's outputs commit: None
|
|
if not sealed:
|
|
sealed.append(True)
|
|
# The writer commits and notifies now; the path is still pending
|
|
self.record_run(
|
|
self.db, self.output, [("", "race_00001_.png")], [LLM_TEXT]
|
|
)
|
|
return result
|
|
|
|
self.handler._history_prompt = history
|
|
self.assertTrue(self.handler.schedule(path))
|
|
self.assertTrue(self.wait_for(lambda: self.linked))
|
|
self.assertTrue(self.wait_for(lambda: self.handler.pending_count() == 0))
|
|
self.assertEqual(self.linked, [("race_00001_.png", prompt_id)])
|
|
|
|
def test_a_seal_while_the_pending_set_is_full_is_rechecked_later(self):
|
|
from utils import image_monitor
|
|
|
|
prompt_id = self.save(LLM_TEXT)
|
|
path = self.image("full_00001_.png")
|
|
self.handler.process_new_image(path) # waiting, unlinked
|
|
blocker = self.image("blocker.png")
|
|
release = threading.Event()
|
|
original = self.handler.process_new_image
|
|
|
|
def slow(image_path, prompt_snapshot=None):
|
|
if image_path == blocker:
|
|
release.wait(5)
|
|
return original(image_path, prompt_snapshot=prompt_snapshot)
|
|
|
|
self.handler.process_new_image = slow
|
|
with mock.patch.object(image_monitor, "MAX_PENDING_IMAGES", 1):
|
|
self.assertTrue(self.handler.schedule(blocker))
|
|
self.assertTrue(self.wait_for(lambda: self.handler.pending_count() == 1))
|
|
self.record_run(self.db, self.output, [("", "full_00001_.png")], [LLM_TEXT])
|
|
self.assertEqual(self.linked, [], "not scheduled while the set is full")
|
|
release.set()
|
|
self.assertTrue(self.wait_for(lambda: self.linked))
|
|
self.assertTrue(self.wait_for(lambda: self.handler.pending_count() == 0))
|
|
self.assertEqual(self.linked, [("full_00001_.png", prompt_id)])
|
|
|
|
def test_a_run_naming_nothing_stored_leaves_the_image_unlinked(self):
|
|
path = self.image("empty_00001_.png")
|
|
self.handler.process_new_image(path)
|
|
self.record_run(self.db, self.output, [("", "empty_00001_.png")], [""])
|
|
self.assertTrue(self.wait_for(lambda: not self.handler.awaiting_history()))
|
|
self.assertTrue(self.wait_for(lambda: self.handler.pending_count() == 0))
|
|
self.assertEqual(self.linked, [])
|
|
|
|
def test_history_found_at_once_leaves_nothing_awaiting(self):
|
|
prompt_id = self.save(LLM_TEXT)
|
|
self.record_run(self.db, self.output, [("", "done_00001_.png")], [LLM_TEXT])
|
|
self.handler.process_new_image(self.image("done_00001_.png"))
|
|
self.assertEqual(self.linked, [("done_00001_.png", prompt_id)])
|
|
self.assertEqual(self.handler.awaiting_history(), [])
|
|
|
|
def test_keyed_images_do_not_wait_for_history(self):
|
|
with mock.patch.object(
|
|
self.handler.metadata_extractor,
|
|
"extract_metadata",
|
|
return_value={"prompt": llm_text_graph()},
|
|
):
|
|
self.handler.process_new_image(self.image("keyed_00001_.png"))
|
|
self.assertEqual(self.handler.awaiting_history(), [])
|
|
|
|
def test_the_waiting_set_is_bounded_and_expires(self):
|
|
from utils import image_monitor
|
|
|
|
with mock.patch.object(image_monitor, "MAX_AWAITING_HISTORY", 2):
|
|
for n in range(3):
|
|
self.handler.process_new_image(self.image(f"b_{n}.png"))
|
|
self.assertEqual(
|
|
[os.path.basename(p) for p in self.handler.awaiting_history()],
|
|
["b_1.png", "b_2.png"],
|
|
)
|
|
later = time.monotonic() + image_monitor.AWAIT_HISTORY_SECONDS + 1
|
|
with mock.patch.object(image_monitor.time, "monotonic", return_value=later):
|
|
self.assertEqual(self.handler.awaiting_history(), [])
|
|
|
|
def test_shutdown_forgets_waiting_images_and_a_dead_handler_is_dropped(self):
|
|
import gc
|
|
|
|
from database import run_attribution
|
|
|
|
self.handler.process_new_image(self.image("gone_00001_.png"))
|
|
self.handler.shutdown()
|
|
self.assertEqual(self.handler.awaiting_history(), [])
|
|
other = ImageGenerationHandler(self.db, make_tracker(self.db))
|
|
with mock.patch.object(
|
|
other.metadata_extractor, "extract_metadata", return_value=None
|
|
):
|
|
other.process_new_image(self.image("gone_00002_.png"))
|
|
del other
|
|
gc.collect()
|
|
run_attribution.notify_outputs_sealed("run", []) # no dead callback raises
|
|
|
|
|
|
def folder_is_case_sensitive(folder):
|
|
"""True when ``folder``'s volume tells names apart by case: the existing
|
|
folder spelled with swapped case is not found (#476)."""
|
|
head, name = os.path.split(folder)
|
|
return not os.path.exists(os.path.join(head, name.swapcase()))
|
|
|
|
|
|
class TestMonitorForeignRuns(RecordedRuns, unittest.TestCase):
|
|
"""#247: outputs of workflows without a PromptManager node, in the monitor."""
|
|
|
|
def setUp(self):
|
|
history_tests.HistoryDatabaseTests.setUp(self)
|
|
self.db = history_tests.HistoryDatabaseTests.database(self)
|
|
self.output = str(Path(self.temp.name) / "output")
|
|
os.makedirs(os.path.join(self.output, "dataset"))
|
|
self.tracker = make_tracker(self.db)
|
|
self.handler = ImageGenerationHandler(self.db, self.tracker)
|
|
self.handler.processing_delay = 0
|
|
self.handler.settle_interval = 0.01
|
|
self.addCleanup(close_handler, self, self.handler)
|
|
self.metadata = {"prompt": shapes.shared_prefix()}
|
|
saved = GalleryConfig.SAVE_FOREIGN_RUNS
|
|
self.addCleanup(setattr, GalleryConfig, "SAVE_FOREIGN_RUNS", saved)
|
|
GalleryConfig.SAVE_FOREIGN_RUNS = True
|
|
patches = (
|
|
mock.patch(
|
|
"utils.image_monitor.host_output_directory", return_value=self.output
|
|
),
|
|
mock.patch(
|
|
"database.run_attribution.host_output_directory",
|
|
return_value=self.output,
|
|
),
|
|
mock.patch.object(
|
|
self.handler.metadata_extractor,
|
|
"extract_metadata",
|
|
side_effect=lambda path: self.metadata,
|
|
),
|
|
)
|
|
for patch in patches:
|
|
patch.start()
|
|
self.addCleanup(patch.stop)
|
|
|
|
close_databases = history_tests.HistoryDatabaseTests.close_databases
|
|
wait_for = TestMonitorHistoryRecheck.wait_for
|
|
|
|
def image(self, name, subfolder="dataset"):
|
|
path = os.path.join(self.output, subfolder, name)
|
|
Path(path).write_bytes(b"png")
|
|
return path
|
|
|
|
def linked_text(self, path):
|
|
row = (
|
|
self.db.model.get_connection()
|
|
.execute(
|
|
"SELECT p.text FROM generated_images g "
|
|
"JOIN prompts p ON p.id = g.prompt_id WHERE g.image_path = ?",
|
|
(str(path),),
|
|
)
|
|
.fetchone()
|
|
)
|
|
return row[0] if row else None
|
|
|
|
def key(self, path):
|
|
return self.handler._awaiting_key(path)
|
|
|
|
def prompt_count(self):
|
|
connection = self.db.model.get_connection()
|
|
return connection.execute("SELECT COUNT(*) FROM prompts").fetchone()[0]
|
|
|
|
def settled(self):
|
|
return self.wait_for(lambda: self.handler.pending_count() == 0)
|
|
|
|
def test_each_dataset_file_links_to_its_branch_once_the_run_ends(self):
|
|
"""CUJ 1: unlinked while the run is open, its own branch after the seal."""
|
|
paths = [self.image(f"IMG_GEN_0000{n}_.png") for n in (1, 2)]
|
|
for path in paths:
|
|
self.handler.process_new_image(path)
|
|
self.assertEqual([self.linked_text(path) for path in paths], [None, None])
|
|
self.assertEqual(
|
|
self.handler.awaiting_foreign(), [self.key(path) for path in paths]
|
|
)
|
|
self.assertEqual(self.handler.awaiting_history(), [])
|
|
self.record_run(
|
|
self.db,
|
|
self.output,
|
|
[("dataset", os.path.basename(path)) for path in paths],
|
|
[],
|
|
nodes=["13", "23"],
|
|
)
|
|
self.assertTrue(self.wait_for(lambda: all(map(self.linked_text, paths))))
|
|
self.assertTrue(self.settled())
|
|
self.assertEqual(
|
|
[self.linked_text(path) for path in paths], list(shapes.BRANCH_TEXTS[:2])
|
|
)
|
|
self.assertEqual(self.handler.awaiting_foreign(), [])
|
|
prompt = self.db.get_prompt_by_hash(
|
|
generate_prompt_hash(shapes.BRANCH_TEXTS[0])
|
|
)
|
|
self.assertEqual(prompt["category"], "scanned")
|
|
self.assertIn("Saved from a workflow run: IMG_GEN_00001_.png", prompt["notes"])
|
|
|
|
def test_a_run_sealed_before_the_file_settles_links_it_at_once(self):
|
|
path = self.image("IMG_GEN_00003_.png")
|
|
self.record_run(
|
|
self.db, self.output, [("dataset", "IMG_GEN_00003_.png")], [], nodes=["33"]
|
|
)
|
|
self.handler.process_new_image(path)
|
|
self.assertEqual(self.linked_text(path), shapes.BRANCH_TEXTS[2])
|
|
self.assertEqual(self.handler.awaiting_foreign(), [])
|
|
|
|
def test_a_name_two_runs_reported_stays_unlinked(self):
|
|
"""CUJ 3: the file was overwritten; no run proves its saver."""
|
|
path = self.image("IMG_GEN_00004_.png")
|
|
reported = [("dataset", "IMG_GEN_00004_.png")]
|
|
self.record_run(self.db, self.output, reported, [], nodes=["13"])
|
|
self.record_run(self.db, self.output, reported, [], item=1, nodes=["23"])
|
|
self.handler.process_new_image(path)
|
|
self.assertIsNone(self.linked_text(path))
|
|
reference = {"folder_type": "output", "subfolder": "dataset"}
|
|
reference["filename"] = "IMG_GEN_00004_.png"
|
|
with self.assertLogs("prompt_manager.image_monitor", "INFO") as logs:
|
|
self.handler._recheck_foreign(path, reference)
|
|
self.assertIn(
|
|
"No certain prompt for IMG_GEN_00004_.png after its run ended; "
|
|
"left unlinked",
|
|
"\n".join(logs.output),
|
|
)
|
|
self.assertIsNone(self.linked_text(path))
|
|
self.assertEqual(self.prompt_count(), 0)
|
|
|
|
def test_the_setting_off_stores_nothing_at_settle_or_after_the_run(self):
|
|
"""CUJ 4: read for every file, so it applies without a restart."""
|
|
GalleryConfig.SAVE_FOREIGN_RUNS = False
|
|
skipped = self.image("IMG_GEN_00005_.png")
|
|
with self.assertLogs("prompt_manager.image_monitor", "INFO") as logs:
|
|
self.handler.process_new_image(skipped)
|
|
self.assertIn(
|
|
"No PromptManager prompt in workflow, skipping image",
|
|
"\n".join(logs.output),
|
|
)
|
|
self.assertEqual(self.handler.awaiting_foreign(), [])
|
|
GalleryConfig.SAVE_FOREIGN_RUNS = True
|
|
waiting = self.image("IMG_GEN_00006_.png")
|
|
self.handler.process_new_image(waiting)
|
|
self.assertEqual(self.handler.awaiting_foreign(), [self.key(waiting)])
|
|
GalleryConfig.SAVE_FOREIGN_RUNS = False
|
|
self.record_run(
|
|
self.db, self.output, [("dataset", "IMG_GEN_00006_.png")], [], nodes=["13"]
|
|
)
|
|
self.assertTrue(self.wait_for(lambda: not self.handler.awaiting_foreign()))
|
|
self.assertTrue(self.settled())
|
|
self.assertEqual(
|
|
[self.linked_text(p) for p in (skipped, waiting)], [None, None]
|
|
)
|
|
self.assertEqual(self.prompt_count(), 0)
|
|
|
|
def test_a_plain_foreign_workflow_links_at_settle_and_counts_no_run(self):
|
|
"""CUJ 2: the default profile decides; the tracker queue is untouched."""
|
|
text = "a lighthouse on a cliff at dusk"
|
|
self.metadata = {"prompt": sampled_graph(text, "blurry")}
|
|
queued = self.db.save_prompt(
|
|
text="queued item", prompt_hash=generate_prompt_hash("queued item")
|
|
)
|
|
self.tracker.set_current_prompt("queued item", {"prompt_id": queued})
|
|
first, second = self.image("plain_00001_.png"), self.image("plain_00002_.png")
|
|
self.handler.process_new_image(first)
|
|
self.handler.process_new_image(second)
|
|
self.assertEqual(
|
|
[self.linked_text(first), self.linked_text(second)], [text, text]
|
|
)
|
|
prompt = self.db.get_prompt_by_hash(generate_prompt_hash(text))
|
|
self.assertEqual(prompt["run_count"] or 0, 0)
|
|
self.assertIsNone(prompt.get("last_used_at"))
|
|
self.assertEqual(self.tracker.pop_next_prompt()["id"], queued)
|
|
self.assertEqual(self.handler.awaiting_foreign(), [])
|
|
|
|
def test_the_sealed_reference_spells_the_file_as_history_recorded_it(self):
|
|
"""Windows: the watcher's spelling differs in case from the run's report.
|
|
|
|
On a case-sensitive folder the two spellings are two files and the
|
|
folded key joins them; on a case-insensitive one (Windows, default
|
|
macOS) they already name one file, so only history's spelling and
|
|
the watcher's differ (#476)."""
|
|
seen = os.path.join(self.output, "dataset", "img_gen_00007_.png")
|
|
Path(seen).write_bytes(b"png")
|
|
recorded = os.path.join(self.output, "Dataset", "IMG_GEN_00007_.png")
|
|
self.assertNotEqual(os.path.basename(recorded), os.path.basename(seen))
|
|
if folder_is_case_sensitive(os.path.join(self.output, "dataset")):
|
|
Path(self.output, "Dataset").mkdir()
|
|
self.image("IMG_GEN_00007_.png", subfolder="Dataset")
|
|
else:
|
|
self.assertTrue(os.path.exists(recorded))
|
|
folded = staticmethod(lambda path: os.path.realpath(path).lower())
|
|
with mock.patch.object(ImageGenerationHandler, "_awaiting_key", folded):
|
|
self.handler.process_new_image(seen)
|
|
self.assertIsNone(self.linked_text(seen))
|
|
self.record_run(
|
|
self.db,
|
|
self.output,
|
|
[("Dataset", os.path.basename(recorded))],
|
|
[],
|
|
nodes=["43"],
|
|
)
|
|
self.assertTrue(self.wait_for(lambda: self.linked_text(seen)))
|
|
self.assertTrue(self.settled())
|
|
self.assertEqual(self.linked_text(seen), shapes.BRANCH_TEXTS[3])
|
|
|
|
def test_a_full_foreign_set_never_evicts_a_waiting_keyless_file(self):
|
|
from utils import image_monitor
|
|
|
|
keyless = self.image("bare_00001_.png", subfolder="")
|
|
with mock.patch.object(
|
|
self.handler.metadata_extractor, "extract_metadata", return_value=None
|
|
):
|
|
self.handler.process_new_image(keyless)
|
|
paths = [self.image(f"IMG_GEN_1000{n}_.png") for n in range(3)]
|
|
with (
|
|
mock.patch.object(image_monitor, "MAX_AWAITING_HISTORY", 2),
|
|
self.assertLogs("prompt_manager.image_monitor", "INFO") as logs,
|
|
):
|
|
for path in paths:
|
|
self.handler.process_new_image(path)
|
|
self.assertEqual(
|
|
self.handler.awaiting_foreign(), [self.key(p) for p in paths[1:]]
|
|
)
|
|
self.assertEqual(self.handler.awaiting_history(), [self.key(keyless)])
|
|
self.assertIn(
|
|
"Too many files wait for their run; IMG_GEN_10000_.png left unlinked",
|
|
"\n".join(logs.output),
|
|
)
|
|
|
|
def test_an_expired_foreign_entry_is_logged_and_left_unlinked(self):
|
|
from utils import image_monitor
|
|
|
|
path = self.image("IMG_GEN_00008_.png")
|
|
self.handler.process_new_image(path)
|
|
later = time.monotonic() + image_monitor.AWAIT_HISTORY_SECONDS + 1
|
|
with (
|
|
mock.patch.object(image_monitor.time, "monotonic", return_value=later),
|
|
self.assertLogs("prompt_manager.image_monitor", "INFO") as logs,
|
|
):
|
|
self.assertEqual(self.handler.awaiting_foreign(), [])
|
|
self.assertIn(
|
|
"No recorded run for IMG_GEN_00008_.png; left unlinked",
|
|
"\n".join(logs.output),
|
|
)
|
|
self.assertIsNone(self.linked_text(path))
|
|
|
|
def test_a_seal_while_the_pending_set_is_full_is_rechecked_later(self):
|
|
from utils import image_monitor
|
|
|
|
path = self.image("IMG_GEN_00009_.png")
|
|
self.handler.process_new_image(path)
|
|
blocker = self.image("blocker.png", subfolder="")
|
|
release = threading.Event()
|
|
original = self.handler.process_new_image
|
|
|
|
def slow(image_path, prompt_snapshot=None):
|
|
if image_path == blocker:
|
|
release.wait(5)
|
|
return None
|
|
return original(image_path, prompt_snapshot=prompt_snapshot)
|
|
|
|
self.handler.process_new_image = slow
|
|
with mock.patch.object(image_monitor, "MAX_PENDING_IMAGES", 1):
|
|
self.assertTrue(self.handler.schedule(blocker))
|
|
self.assertTrue(self.wait_for(lambda: self.handler.pending_count() == 1))
|
|
self.record_run(
|
|
self.db,
|
|
self.output,
|
|
[("dataset", "IMG_GEN_00009_.png")],
|
|
[],
|
|
nodes=["13"],
|
|
)
|
|
self.assertEqual(list(self.handler._foreign_due), [path])
|
|
release.set()
|
|
self.assertTrue(self.wait_for(lambda: self.linked_text(path)))
|
|
self.assertTrue(self.settled())
|
|
self.assertEqual(self.linked_text(path), shapes.BRANCH_TEXTS[0])
|
|
|
|
def test_shutdown_forgets_waiting_and_due_foreign_files(self):
|
|
path = self.image("IMG_GEN_00010_.png")
|
|
self.handler.process_new_image(path)
|
|
self.handler._foreign_due[path] = {"filename": "IMG_GEN_00010_.png"}
|
|
self.handler.shutdown()
|
|
self.assertEqual(self.handler.awaiting_foreign(), [])
|
|
self.assertEqual(list(self.handler._foreign_due), [])
|
|
|
|
def link_count(self, path):
|
|
connection = self.db.model.get_connection()
|
|
return connection.execute(
|
|
"SELECT COUNT(*) FROM generated_images WHERE image_path = ?", (str(path),)
|
|
).fetchone()[0]
|
|
|
|
def test_a_file_linked_to_another_prompt_never_gets_a_second_link(self):
|
|
"""One prompt per file: the pair path and the positive path both refuse."""
|
|
other = self.db.save_prompt(
|
|
"an earlier prompt", generate_prompt_hash("an earlier prompt")
|
|
)
|
|
plain = self.image("plain_00004_.png")
|
|
self.db.link_image_to_prompt(other, plain)
|
|
self.metadata = {"prompt": sampled_graph("a quiet harbor", "blurry")}
|
|
with self.assertLogs("prompt_manager.image_monitor", "INFO") as logs:
|
|
self.handler.process_new_image(plain)
|
|
self.assertIn("Foreign run: conflict", "\n".join(logs.output))
|
|
dataset = self.image("IMG_GEN_00011_.png")
|
|
self.db.link_image_to_prompt(other, dataset)
|
|
self.metadata = {"prompt": shapes.shared_prefix()}
|
|
self.record_run(
|
|
self.db, self.output, [("dataset", "IMG_GEN_00011_.png")], [], nodes=["13"]
|
|
)
|
|
self.handler.process_new_image(dataset)
|
|
for path in (plain, dataset):
|
|
self.assertEqual(self.linked_text(path), "an earlier prompt")
|
|
self.assertEqual(self.link_count(path), 1)
|
|
for text in ("a quiet harbor", shapes.BRANCH_TEXTS[0]):
|
|
self.assertIsNone(self.db.get_prompt_by_hash(generate_prompt_hash(text)))
|
|
|
|
def test_the_file_size_is_read_only_when_the_ladder_needs_it(self):
|
|
"""PR 1's lazy size: a foreign file decided without a store reads none."""
|
|
# Foreign (a core encoder only) and no API graph: never, at step 3
|
|
encoder = {"type": "CLIPTextEncode", "inputs": [], "widgets_values": ["x"]}
|
|
self.metadata = {"text_encoder_nodes": [encoder]}
|
|
path = self.image("bare_00002_.png", subfolder="")
|
|
with mock.patch("utils.foreign_runs.os.path.getsize") as size:
|
|
self.handler.process_new_image(path)
|
|
size.assert_not_called()
|
|
|
|
def untraceable(self):
|
|
"""Unresolved and only the legacy parser names a text (an unknown node)."""
|
|
graph = sampled_graph("unused", "blurry")
|
|
graph["1"]["inputs"]["text"] = ["30", 0]
|
|
graph["30"] = {
|
|
"class_type": "VendorPromptMaker",
|
|
"inputs": {"prompt": "a misty valley at dawn"},
|
|
}
|
|
return {"prompt": graph}
|
|
|
|
def test_no_record_at_settle_keeps_the_file_waiting_for_the_legacy_floor(self):
|
|
"""Review: a lookup that found no record is not the run's answer yet."""
|
|
self.metadata = self.untraceable()
|
|
path = self.image("vendor_00001_.png")
|
|
self.handler.process_new_image(path)
|
|
self.assertIsNone(self.linked_text(path))
|
|
self.assertEqual(self.handler.awaiting_foreign(), [self.key(path)])
|
|
self.record_run(self.db, self.output, [("dataset", "vendor_00001_.png")], [])
|
|
self.assertTrue(self.wait_for(lambda: self.linked_text(path)))
|
|
self.assertTrue(self.settled())
|
|
self.assertEqual(self.linked_text(path), "a misty valley at dawn")
|
|
|
|
def test_a_record_found_at_settle_without_a_certain_prompt_ends_the_wait(self):
|
|
"""Review: sealed before settle, still ambiguous: no seal is coming."""
|
|
path = self.image("IMG_GEN_00012_.png")
|
|
# One run reported it from two savers: a record, but no saver
|
|
reported = [("dataset", "IMG_GEN_00012_.png")] * 2
|
|
self.record_run(self.db, self.output, reported, [], nodes=["13", "23"])
|
|
with self.assertLogs("prompt_manager.image_monitor", "INFO") as logs:
|
|
self.handler.process_new_image(path)
|
|
self.assertIn(
|
|
"No certain prompt for IMG_GEN_00012_.png; left unlinked",
|
|
"\n".join(logs.output),
|
|
)
|
|
self.assertEqual(self.handler.awaiting_foreign(), [])
|
|
self.assertIsNone(self.linked_text(path))
|
|
|
|
def test_monitoring_saved_off_stops_foreign_writes_at_settle_and_recheck(self):
|
|
"""Review: the observer runs until a restart; nothing is written."""
|
|
enabled = GalleryConfig.MONITORING_ENABLED
|
|
self.addCleanup(setattr, GalleryConfig, "MONITORING_ENABLED", enabled)
|
|
GalleryConfig.MONITORING_ENABLED = False
|
|
self.metadata = {"prompt": sampled_graph("a quiet harbor", "blurry")}
|
|
plain = self.image("plain_00005_.png")
|
|
self.handler.process_new_image(plain)
|
|
self.assertIsNone(self.linked_text(plain))
|
|
GalleryConfig.MONITORING_ENABLED = True
|
|
self.metadata = {"prompt": shapes.shared_prefix()}
|
|
waiting = self.image("IMG_GEN_00013_.png")
|
|
self.handler.process_new_image(waiting)
|
|
GalleryConfig.MONITORING_ENABLED = False
|
|
reference = {"folder_type": "output", "subfolder": "dataset"}
|
|
reference["filename"] = "IMG_GEN_00013_.png"
|
|
self.record_run(
|
|
self.db, self.output, [("dataset", reference["filename"])], [], nodes=["13"]
|
|
)
|
|
self.assertTrue(self.wait_for(lambda: not self.handler.awaiting_foreign()))
|
|
self.assertTrue(self.settled())
|
|
self.handler._recheck_foreign(waiting, reference)
|
|
self.assertIsNone(self.linked_text(waiting))
|
|
self.assertEqual(self.prompt_count(), 0)
|
|
|
|
def test_the_setting_is_read_again_just_before_the_store(self):
|
|
self.metadata = {"prompt": sampled_graph("a quiet harbor", "blurry")}
|
|
path = self.image("plain_00006_.png")
|
|
real = __import__("utils.foreign_runs", fromlist=["x"]).decide_foreign
|
|
|
|
def decide_then_switch_off(*args, **kwargs):
|
|
decision = real(*args, **kwargs)
|
|
GalleryConfig.SAVE_FOREIGN_RUNS = False
|
|
return decision
|
|
|
|
with mock.patch(
|
|
"utils.image_monitor.decide_foreign", side_effect=decide_then_switch_off
|
|
):
|
|
self.handler.process_new_image(path)
|
|
self.assertIsNone(self.linked_text(path))
|
|
self.assertEqual(self.prompt_count(), 0)
|
|
|
|
def test_a_seal_between_registering_and_the_settle_lookup_links_once(self):
|
|
"""Regression guard (passes before the review fixes): one link only."""
|
|
path = self.image("IMG_GEN_00014_.png")
|
|
real = __import__("utils.foreign_runs", fromlist=["x"]).history_lookup
|
|
|
|
def sealed_first(db, image_path, reference=None):
|
|
lookup = real(db, image_path, reference)
|
|
|
|
def run_then_look():
|
|
if reference is None: # the settle lookup: the run seals now
|
|
self.record_run(
|
|
self.db,
|
|
self.output,
|
|
[("dataset", "IMG_GEN_00014_.png")],
|
|
[],
|
|
nodes=["13"],
|
|
)
|
|
return lookup()
|
|
|
|
return run_then_look
|
|
|
|
with mock.patch("utils.image_monitor.history_lookup", side_effect=sealed_first):
|
|
self.assertTrue(self.handler.schedule(path))
|
|
self.assertTrue(self.wait_for(lambda: self.linked_text(path)))
|
|
self.assertTrue(self.settled())
|
|
self.assertEqual(self.linked_text(path), shapes.BRANCH_TEXTS[0])
|
|
self.assertEqual(self.link_count(path), 1)
|
|
self.assertEqual(list(self.handler._foreign_due), [])
|
|
|
|
def test_a_drained_recheck_never_drops_a_newer_reference(self):
|
|
path = self.image("IMG_GEN_00015_.png")
|
|
older, newer = {"filename": "old"}, {"filename": "new"}
|
|
self.handler._foreign_due[path] = older
|
|
real = self.handler.schedule
|
|
|
|
def schedule(image_path, snapshot=None):
|
|
# The run is sealed again while the older one is scheduled
|
|
self.handler._foreign_due[image_path] = newer
|
|
return real(image_path, snapshot)
|
|
|
|
self.handler.schedule = schedule
|
|
with mock.patch.object(self.handler, "process_new_image"):
|
|
self.handler._drain_history_due()
|
|
self.assertIs(self.handler._foreign_due[path], newer)
|
|
|
|
def test_an_evicted_due_recheck_is_logged(self):
|
|
from utils import image_monitor
|
|
|
|
self.handler.schedule = lambda *args, **kwargs: False
|
|
with (
|
|
mock.patch.object(image_monitor, "MAX_AWAITING_HISTORY", 1),
|
|
self.assertLogs("prompt_manager.image_monitor", "INFO") as logs,
|
|
):
|
|
for name in ("a_00001_.png", "b_00001_.png"):
|
|
self.handler._schedule_recheck(self.image(name), {"filename": name})
|
|
self.assertIn(
|
|
"Too many files wait for their run; a_00001_.png left unlinked",
|
|
"\n".join(logs.output),
|
|
)
|
|
|
|
def test_a_failed_store_is_logged_and_leaves_the_file_unlinked(self):
|
|
self.metadata = {"prompt": sampled_graph("a quiet harbor", "blurry")}
|
|
path = self.image("plain_00003_.png")
|
|
with (
|
|
mock.patch(
|
|
"utils.foreign_runs.ingest_resolved_pair",
|
|
side_effect=sqlite3.OperationalError("locked"),
|
|
),
|
|
self.assertLogs("prompt_manager.image_monitor", "ERROR"),
|
|
):
|
|
self.handler.process_new_image(path)
|
|
self.assertIsNone(self.linked_text(path))
|
|
|
|
|
|
class TestBoundedLookup(RecordedRuns, unittest.TestCase):
|
|
"""M5: history lookups are batched, and the history schema is unchanged."""
|
|
|
|
setUp = history_tests.HistoryDatabaseTests.setUp
|
|
close_databases = history_tests.HistoryDatabaseTests.close_databases
|
|
database = history_tests.HistoryDatabaseTests.database
|
|
|
|
def output_lookups(self, db, action):
|
|
statements = []
|
|
conn = db.model.get_connection()
|
|
conn.set_trace_callback(statements.append)
|
|
try:
|
|
result = action()
|
|
finally:
|
|
conn.set_trace_callback(None)
|
|
return result, [s for s in statements if "FROM execution_run_outputs" in s]
|
|
|
|
def test_one_query_per_chunk_not_per_file_and_same_answers(self):
|
|
from database import run_attribution
|
|
|
|
db = self.database()
|
|
output = str(Path(self.temp.name) / "output")
|
|
os.makedirs(output)
|
|
prompt_id = db.save_prompt(
|
|
text=LLM_TEXT, prompt_hash=generate_prompt_hash(LLM_TEXT)
|
|
)
|
|
names = [f"img_{n:05}_.png" for n in range(25)]
|
|
paths = [os.path.join(output, name) for name in names]
|
|
for path in paths:
|
|
Path(path).write_bytes(b"png")
|
|
self.record_run(db, output, [("", n) for n in names[:10]], [LLM_TEXT])
|
|
self.record_run(db, output, [("", n) for n in names[5:10]], [""], item=1)
|
|
with mock.patch.object(run_attribution, "LOOKUP_CHUNK", 10):
|
|
result, lookups = self.output_lookups(
|
|
db, lambda: run_attribution.stored_run_prompts(db, paths, output)
|
|
)
|
|
self.assertEqual(len(lookups), 3) # 25 names in chunks of 10
|
|
expected = {path: None for path in paths}
|
|
expected.update({path: prompt_id for path in paths[:5]})
|
|
self.assertEqual(result, expected)
|
|
for path in paths[:12]:
|
|
self.assertEqual(
|
|
run_attribution.stored_run_prompt(db, path, output), expected[path]
|
|
)
|
|
|
|
def test_the_history_schema_gains_no_index(self):
|
|
from database.execution_history_schema import history_schema_error
|
|
|
|
db = self.database()
|
|
conn = db.model.get_connection()
|
|
indexes = [
|
|
row[1]
|
|
for row in conn.execute("PRAGMA index_list(execution_run_outputs)")
|
|
if row[3] != "pk"
|
|
]
|
|
self.assertEqual(indexes, [])
|
|
self.assertIsNone(history_schema_error(conn, allow_missing=False))
|
|
|
|
|
|
class TestScanBatchLookup(RecordedRuns, AdminAPITestCase):
|
|
"""M5: a scan batch asks history once per chunk of names, not per file."""
|
|
|
|
def test_a_batch_its_metadata_links_never_asks_history(self):
|
|
"""R3: the batch lookup is made only when a file needs it.
|
|
|
|
Since #247 every file the default profile leaves undecided asks (the
|
|
recorded saver may name its branch), so these files resolve as pairs.
|
|
"""
|
|
files = []
|
|
for n in range(3):
|
|
path = self.output_dir / f"typed_{n:05}_.png"
|
|
path.write_bytes(b"png")
|
|
files.append(path)
|
|
meta = {"prompt": json.dumps(sampled_graph("typed mountains"))}
|
|
statements = []
|
|
conn = self.api.db.model.get_connection()
|
|
conn.set_trace_callback(statements.append)
|
|
try:
|
|
counts = self.api._ingest_scan_batch_sync([(f, meta) for f in files])
|
|
finally:
|
|
conn.set_trace_callback(None)
|
|
self.assertEqual(counts["added"] + counts["linked"], 3, counts)
|
|
lookups = [s for s in statements if "FROM execution_run_outputs" in s]
|
|
self.assertEqual(lookups, [])
|
|
|
|
def test_metadata_less_batch_links_with_one_history_lookup(self):
|
|
db = self.api.db
|
|
prompt_id = db.save_prompt(
|
|
text=LLM_TEXT, prompt_hash=generate_prompt_hash(LLM_TEXT)
|
|
)
|
|
files = []
|
|
for n in range(6):
|
|
path = self.output_dir / f"bare_{n:05}_.png"
|
|
path.write_bytes(b"png")
|
|
files.append(path)
|
|
self.record_run(
|
|
db, str(self.output_dir), [("", f.name) for f in files[:4]], [LLM_TEXT]
|
|
)
|
|
statements = []
|
|
conn = db.model.get_connection()
|
|
conn.set_trace_callback(statements.append)
|
|
try:
|
|
counts = self.api._ingest_scan_batch_sync([(f, {}) for f in files])
|
|
finally:
|
|
conn.set_trace_callback(None)
|
|
lookups = [s for s in statements if "FROM execution_run_outputs" in s]
|
|
self.assertEqual(len(lookups), 1)
|
|
self.assertEqual(counts["linked"], 4)
|
|
linked = {
|
|
row[0]
|
|
for row in conn.execute(
|
|
"SELECT filename FROM generated_images WHERE prompt_id=?",
|
|
(prompt_id,),
|
|
)
|
|
}
|
|
self.assertEqual(linked, {f.name for f in files[:4]})
|
|
|
|
def test_the_batch_saver_equals_the_files_recorded_output(self):
|
|
from py.api.admin import _BatchHistory
|
|
|
|
db = self.api.db
|
|
names = ["one_node.png", "two_nodes.png", "two_runs.png", "unrecorded.png"]
|
|
files = []
|
|
for name in names:
|
|
path = self.output_dir / name
|
|
path.write_bytes(b"png")
|
|
files.append(path)
|
|
root = str(self.output_dir)
|
|
self.record_run(db, root, [("", "one_node.png")], [], nodes=["13"])
|
|
self.record_run(
|
|
db, root, [("", "two_nodes.png")] * 2, [], item=1, nodes=["13", "23"]
|
|
)
|
|
self.record_run(db, root, [("", "two_runs.png")], [], item=2)
|
|
self.record_run(db, root, [("", "two_runs.png")], [], item=3)
|
|
history = _BatchHistory(self.api, files)
|
|
with mock.patch(
|
|
"database.run_attribution.host_output_directory", return_value=root
|
|
):
|
|
for path in files:
|
|
with self.subTest(path=path.name):
|
|
expected = recorded_output(db, str(path), root)
|
|
found = history.recorded(str(path))
|
|
if expected is None:
|
|
self.assertIsNone(found)
|
|
else:
|
|
self.assertEqual(found.saver_id, expected.saver_id)
|
|
self.assertEqual(history.recorded(str(files[0])).saver_id, "13")
|
|
self.assertIsNone(history.recorded(str(files[1])).saver_id)
|
|
|
|
def test_a_failed_batch_lookup_is_logged_and_names_nothing(self):
|
|
from py.api.admin import _BatchHistory
|
|
|
|
path = self.output_dir / "x.png"
|
|
history = _BatchHistory(self.api, [path])
|
|
with (
|
|
mock.patch(
|
|
"py.api.admin.recorded_runs", side_effect=sqlite3.OperationalError("x")
|
|
),
|
|
self.assertLogs("prompt_manager", "WARNING"),
|
|
):
|
|
self.assertIsNone(history.recorded(str(path)))
|
|
self.assertIsNone(history.get(str(path)))
|
|
|
|
def test_a_failed_prompt_lookup_is_logged_and_names_nothing(self):
|
|
from py.api.admin import _BatchHistory
|
|
|
|
path = self.output_dir / "one.png"
|
|
path.write_bytes(b"png")
|
|
root = str(self.output_dir)
|
|
self.record_run(self.api.db, root, [("", path.name)], [LLM_TEXT])
|
|
history = _BatchHistory(self.api, [path])
|
|
with (
|
|
mock.patch(
|
|
"database.run_attribution.host_output_directory", return_value=root
|
|
),
|
|
mock.patch(
|
|
"py.api.admin.named_prompt", side_effect=sqlite3.OperationalError("x")
|
|
),
|
|
self.assertLogs("prompt_manager", "WARNING"),
|
|
):
|
|
self.assertIsNone(history.get(str(path)))
|
|
self.assertEqual(history.recorded(str(path)).saver_id, "9")
|
|
|
|
|
|
class TestLinkedOutputRoot(RecordedRuns, unittest.TestCase):
|
|
"""R2: ComfyUI's output folder reached through a symlink or junction."""
|
|
|
|
setUp = history_tests.HistoryDatabaseTests.setUp
|
|
close_databases = history_tests.HistoryDatabaseTests.close_databases
|
|
database = history_tests.HistoryDatabaseTests.database
|
|
|
|
def linked_root(self):
|
|
self.real = str(Path(self.temp.name) / "big_drive" / "output")
|
|
os.makedirs(os.path.join(self.real, "sub"))
|
|
link = str(Path(self.temp.name) / "output_link")
|
|
try:
|
|
os.symlink(self.real, link, target_is_directory=True)
|
|
except (OSError, NotImplementedError) as error:
|
|
self.skipTest(f"symlinks unavailable: {error}")
|
|
return link
|
|
|
|
def test_either_spelling_of_the_file_is_attributed(self):
|
|
db = self.database()
|
|
link = self.linked_root()
|
|
prompt_id = db.save_prompt(
|
|
text=LLM_TEXT, prompt_hash=generate_prompt_hash(LLM_TEXT)
|
|
)
|
|
real_file = os.path.join(self.real, "sub", "img_00001_.png")
|
|
Path(real_file).write_bytes(b"png")
|
|
# History keeps the host's spelling of its output folder (the link)
|
|
self.record_run(db, link, [("sub", "img_00001_.png")], [LLM_TEXT])
|
|
link_file = os.path.join(link, "sub", "img_00001_.png")
|
|
self.assertEqual(stored_run_prompt(db, real_file, link), prompt_id)
|
|
self.assertEqual(stored_run_prompt(db, link_file, link), prompt_id)
|
|
# The real folder named as the root is another root key: not attributed
|
|
self.assertIsNone(stored_run_prompt(db, real_file, self.real))
|
|
|
|
def test_the_monitor_rechecks_a_file_seen_through_the_real_path(self):
|
|
db = self.database()
|
|
link = self.linked_root()
|
|
prompt_id = db.save_prompt(
|
|
text=LLM_TEXT, prompt_hash=generate_prompt_hash(LLM_TEXT)
|
|
)
|
|
real_file = os.path.join(self.real, "late_00001_.png")
|
|
Path(real_file).write_bytes(b"png")
|
|
handler = ImageGenerationHandler(db, make_tracker(db))
|
|
handler.processing_delay = 0
|
|
handler.settle_interval = 0.01
|
|
self.addCleanup(close_handler, self, handler)
|
|
linked = []
|
|
handler.link_image_to_prompt = lambda path, prompt, meta: linked.append(
|
|
prompt["id"]
|
|
)
|
|
with (
|
|
mock.patch("utils.image_monitor.host_output_directory", return_value=link),
|
|
mock.patch.object(
|
|
handler.metadata_extractor, "extract_metadata", return_value=None
|
|
),
|
|
):
|
|
handler.process_new_image(real_file)
|
|
self.assertEqual(linked, [])
|
|
self.record_run(db, link, [("", "late_00001_.png")], [LLM_TEXT])
|
|
deadline = time.monotonic() + 5
|
|
while not linked and time.monotonic() < deadline:
|
|
time.sleep(0.01)
|
|
self.assertEqual(linked, [prompt_id])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|