Vito
39877202b2
* feat(inspector): use recorded run history when the graph alone can't name the prompt (#386) When a file's own graph leaves the prompt open, both inspect routes ask recorded run history about the file (one run reported it and it was written during that run). The save node that run reported it from anchors the walk like a file-name match, so dataset branches that share one saver prefix show their own prompt, labelled "Recorded run". Still open at one sampler whose PromptManager text was recorded at execution (a language model's output), the recorded text shows as "Execution recorded". Results the file resolves alone never change; a failed lookup judges the file alone. * fix(inspector): recorded texts follow history's saver; conflicts anchor nothing (#386) Review follow-ups: the texts PromptManager recorded are looked up at the sampler history's saver leads to, never at a sampler the unanchored walk named; a file name naming a different saver than history anchors nothing; outputs and texts are read in one transaction; "Recorded run" ranks below every workflow-inferred source; unexpected lookup errors log at error level. * fix(history): look up recorded runs by history's own spelling of the file (#386) The history route inspects a case-normalized target on Windows, while history matches file names exactly, so the occurrence's reference names the file. * refactor(inspector): share the inspector's resolve as inspection_resolution inspect_prompt_evidence now calls inspection_resolution, which returns the resolution to use, the anchored result and whether the recorded saver chose it. Ingestion can call the same helper for files the default profile cannot decide (#247). * feat(ingestion): one per-file prompt decision for the scan and the monitor utils/file_prompt.py holds decide_file_prompt, the ladder the output scan runs today plus the inspector's rule (inspection profile, then the recorded saver) before the legacy parser, with a SCAN and a MONITOR policy. Junk text (blank, numbers, LoRA tags only, model file names) is never stored. graph_prompt_text and parse_comfyui_prompt move here; the parser accepts parsed metadata (#247). * feat(ingestion): store a file's prompt in one transaction, one prompt per file database/scanned_prompts.py links a file to the prompt of one text, creating it when new, inside one BEGIN IMMEDIATE transaction: a file linked to another prompt is a conflict, a failed link leaves no prompt behind and two writers store one prompt. pair_ingestion gains file_link_rows, link_conflict and Marks/mark_created; ingest_resolved_pair takes optional marks for the components it creates. recorded_runs and named_prompt in run_attribution become public (#247). * feat(scan): decide each scanned file with the shared ladder (#247) _ingest_scanned_file calls decide_file_prompt (scan policy) and stores the result: pairs through ingest_resolved_pair and positives through store_scanned_prompt, both with the scan marks. Files the default profile leaves undecided now get the inspector's resolve, with the saver recorded history names, before the legacy parser, so each branch of a shared-prefix dataset run links its own prompt. Junk text is never stored, a file linked to another prompt is a conflict instead of a second link, and a failed link leaves no prompt behind. _BatchHistory reads the stored prompt and the recorded saver from one recorded_runs pass. _parse_comfyui_prompt calls the moved parse_comfyui_prompt. TESTING.md documents the decision. * fix(ingestion): keep the file-name saver anchor out of the ladder (#390) The #390 ruling queues a file that only its saver's file name resolves for Needs review, never linking it. inspection_resolution gains name_anchor; the ladder passes False, so such a file falls to the legacy floor in the scan (and to Needs review) and is never stored by the monitor, while the recorded saver still decides. The inspector keeps the name anchor. * fix(ingestion): no joined PromptManager text from the legacy floor; junk keeps its step's status The legacy parser returns prepend + text + append, so its answer is refused when a PromptManager-family node has a prepend or append part (#343). Junk found at the undecided-graph step keeps that step's unlinked status and review, so an ambiguous file is still queued for Needs review. The monitor policy's ambiguous file no longer asks for the recorded run, and the file size may be read lazily. * fix(ingestion): keep the legacy answer that is a PromptManager node's own text The joined-text refusal also refused the #170 parity files whose PromptManager prepend is a link the browser parser does not follow: its answer is then the node's literal text, which is what the node stores. Refuse only an answer that is not such a node's own text. * fix(scan): one prompt per file for recorded-run links; writes survive a library move The recorded run's stored prompt now links in one transaction with the file's link rows, so a file linked to another prompt is a conflict instead of a second link. The scan's stores run through survive_library_switch. A single file asks history once (a one-file _BatchHistory), the scan passes in_output=False and reads the file size only for files with prompt metadata. A conflict row reads "Kept: already linked to a different prompt". TESTING.md records the rules. * fix(ingestion): match a file's link rows under its resolved spelling too The scan stores realpath spellings and the monitor the watcher's raw path, so a file reached through a symlinked or junction output folder could get a second generated_images row. file_link_rows now queries the path and its realpath (each normalized and raw) and dedupes by row id. TESTING.md notes the reverse order is matched only when the raw spelling is the realpath. * fix(ingestion): a switched-off PromptManager node behind a pass-through stores nothing (#90) The #90 graph walk follows only prompt-like inputs, so a PromptManagerText with Save to library off reaching the encoder through another input name (easy showAnything.anything, a custom node's value) was invisible to it, and the inspector's resolve stored its text. The ladder now checks the positive every step names: a resolved or inspected positive whose path passes a node that is off, or a legacy answer that is such a node's own text, makes the file not_saved. * fix(ingestion): any switched-off node upstream of a positive stores nothing (#90) The ladder walks every input upstream of each sampler or guider positive (not negative, model, clip, latent or image inputs) for PromptManager nodes whose Save to library is off or wired from an unreadable node, so the sampler-less step, an off node with a prepend part and a PromptManagerPair behind a pass-through no longer store its text. The legacy-floor guard reads a node's resolved text. Every guard reports not_saved for a switch that is off and not_saved_wired for one it cannot read. * fix(ingestion): the switch walk skips negative inputs by #90's own rule The upstream walk treated only names containing "negative" as negative, so an off node on a neg, neg_text or text_neg input blocked an enabled positive. prompt_graph gains is_negative_input (the _NEGATIVE_KEY rule, also used by follows_input) and the walk uses it. The guards accept a node with null inputs and a side whose provenance is not an object. * fix(ingestion): the switch walk also skips any input named with "negative" #90's rule matches only whole underscore-separated words, so negative2, negativeprompt and NegativePrompt counted as positive inputs and an off node wired only to one of them blocked the file. The walk skips those too; follows_input is unchanged.
1024 lines
42 KiB
Python
1024 lines
42 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
|
|
from unittest import 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 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):
|
|
setUp = history_tests.HistoryDatabaseTests.setUp
|
|
close_databases = history_tests.HistoryDatabaseTests.close_databases
|
|
database = history_tests.HistoryDatabaseTests.database
|
|
|
|
def setUp(self): # noqa: F811 - 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
|
|
|
|
|
|
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()
|