mARTin-B78
084a99af2c
Bind the consistent-image Idempotency-Key to the character request so a reused key cannot return another character's job; add back-to-back and concurrent two-character regression tests. Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
169 lines
8.1 KiB
Python
169 lines
8.1 KiB
Python
"""Per-character isolation of the consistent-image (profile + reference sheet) generation path.
|
|
|
|
Two deliberately dissimilar characters are generated back to back and concurrently; each ComfyUI
|
|
request must carry only its own prompt, seed family, identity image and result.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
import tempfile
|
|
import threading
|
|
import time
|
|
import unittest
|
|
from unittest import mock
|
|
|
|
os.environ.setdefault("CONFIG_DIR", tempfile.mkdtemp(prefix="ttvc-charisol-test-config-"))
|
|
|
|
from fastapi import FastAPI
|
|
from fastapi.testclient import TestClient
|
|
|
|
import routes.conversation as conv
|
|
|
|
DWARF = {
|
|
"name": "Borin", "book": "Isolation Book",
|
|
"sheet": {"race_species": "Zwerg", "gender": "male", "age": "60", "physical": "stocky dwarf with braided red beard",
|
|
"clothing": "heavy forged-iron mail", "items": ["war hammer"]},
|
|
}
|
|
ELF = {
|
|
"name": "Lirael", "book": "Isolation Book",
|
|
"sheet": {"race_species": "Elf", "gender": "female", "age": "30", "physical": "tall slender elf with silver hair",
|
|
"clothing": "light green linen robe", "items": ["longbow"]},
|
|
}
|
|
|
|
|
|
def _payload(char: dict) -> dict:
|
|
sheet = char["sheet"]
|
|
return {
|
|
"name": char["name"], "book": char["book"], "sheet": sheet, "visual_style": "cinematic",
|
|
"book_profile": {},
|
|
"profile_prompt": f"Profile portrait of {char['name']}: {sheet['physical']}",
|
|
"reference_prompt": f"Reference sheet of {char['name']}: {sheet['physical']}",
|
|
}
|
|
|
|
|
|
def _uri(tag: str) -> str:
|
|
return "data:image/png;base64,UE5H" + tag.encode().hex()
|
|
|
|
|
|
class ConsistentImageIsolationTests(unittest.TestCase):
|
|
def setUp(self):
|
|
self.identity_calls: list[dict] = []
|
|
self.sheet_calls: list[dict] = []
|
|
self.lock = threading.Lock()
|
|
self.delay = 0.0
|
|
conv._consistent_image_jobs.clear()
|
|
conv._consistent_image_idempotency.clear()
|
|
|
|
def fake_identity(url, workflow, node, field, out, prompt, seed_key="", progress=None):
|
|
time.sleep(self.delay)
|
|
with self.lock:
|
|
self.identity_calls.append({"prompt": prompt, "seed_key": seed_key})
|
|
return _uri("identity:" + seed_key.split("|")[3])
|
|
|
|
def fake_sheet(url, workflow, identity_image, profile_prompt, front, side, back, scene, props,
|
|
style_key="", progress=None, face_transfer=True):
|
|
time.sleep(self.delay)
|
|
with self.lock:
|
|
self.sheet_calls.append({"identity": identity_image, "profile": profile_prompt, "front": front,
|
|
"side": side, "back": back, "scene": scene, "props": list(props),
|
|
"style_key": style_key})
|
|
owner = style_key.split("|")[3]
|
|
return {"profile": _uri("profile:" + owner), "reference": _uri("reference:" + owner), "_panels": {}}
|
|
|
|
patches = [
|
|
mock.patch.object(conv, "_load_settings", return_value={"comfyui_url": "http://comfy.invalid",
|
|
"image_check_enabled": False}),
|
|
mock.patch.object(conv, "_comfyui_run", side_effect=fake_identity),
|
|
mock.patch.object(conv, "_comfyui_run_consistent_images", side_effect=fake_sheet),
|
|
mock.patch.object(conv, "_persist_consistent_image_artifacts", side_effect=lambda job, result: dict(result)),
|
|
mock.patch.object(conv, "background_job_create"),
|
|
mock.patch.object(conv, "background_job_update"),
|
|
mock.patch.object(conv, "background_job_get", return_value=None),
|
|
]
|
|
for p in patches:
|
|
p.start()
|
|
self.addCleanup(p.stop)
|
|
app = FastAPI()
|
|
app.include_router(conv.router)
|
|
self.client = TestClient(app)
|
|
|
|
def _run(self, char: dict, headers: dict | None = None) -> dict:
|
|
resp = self.client.post("/api/character-generate-consistent-images", json=_payload(char), headers=headers or {})
|
|
self.assertEqual(resp.status_code, 200, resp.text)
|
|
return self._wait(resp.json()["job_id"])
|
|
|
|
def _wait(self, job_id: str) -> dict:
|
|
for _ in range(200):
|
|
status = self.client.get(f"/api/character-generate-consistent-images/status/{job_id}").json()
|
|
if status.get("done"):
|
|
self.assertIsNone(status.get("error"), status)
|
|
return status
|
|
time.sleep(0.05)
|
|
self.fail("job did not finish")
|
|
|
|
def _assert_separate(self, dwarf_result: dict, elf_result: dict) -> None:
|
|
self.assertEqual(len(self.sheet_calls), 2)
|
|
by_owner = {c["style_key"].split("|")[3]: c for c in self.sheet_calls}
|
|
dwarf_call, elf_call = by_owner["Borin"], by_owner["Lirael"]
|
|
# prompts: each request holds only its own character's traits
|
|
for key in ("profile", "front", "side", "back", "scene"):
|
|
self.assertNotIn("Lirael", dwarf_call[key]); self.assertNotIn("elf", dwarf_call[key].lower().replace("self", ""))
|
|
self.assertNotIn("Borin", elf_call[key]); self.assertNotIn("dwarf", elf_call[key].lower())
|
|
self.assertNotIn("longbow", " ".join(dwarf_call["props"]).lower())
|
|
self.assertNotIn("hammer", " ".join(elf_call["props"]).lower())
|
|
# identity reference image and seed family are per-character
|
|
self.assertIn("identity:Borin".encode().hex(), dwarf_call["identity"])
|
|
self.assertIn("identity:Lirael".encode().hex(), elf_call["identity"])
|
|
self.assertNotEqual(dwarf_call["style_key"], elf_call["style_key"])
|
|
self.assertNotEqual(dwarf_call["identity"], elf_call["identity"])
|
|
for call in self.identity_calls:
|
|
owner = call["seed_key"].split("|")[3]
|
|
other = "Lirael" if owner == "Borin" else "Borin"
|
|
self.assertNotIn(other, call["prompt"])
|
|
# results
|
|
self.assertEqual(dwarf_result["result"]["profile"], _uri("profile:Borin"))
|
|
self.assertEqual(dwarf_result["result"]["reference"], _uri("reference:Borin"))
|
|
self.assertEqual(elf_result["result"]["profile"], _uri("profile:Lirael"))
|
|
self.assertEqual(elf_result["result"]["reference"], _uri("reference:Lirael"))
|
|
self.assertNotEqual(dwarf_result["result"]["identity_fingerprint"], elf_result["result"]["identity_fingerprint"])
|
|
|
|
def test_back_to_back_runs_do_not_share_prompts_refs_or_results(self):
|
|
first = self._run(DWARF)
|
|
second = self._run(ELF)
|
|
self._assert_separate(first, second)
|
|
|
|
def test_concurrent_runs_do_not_share_prompts_refs_or_results(self):
|
|
self.delay = 0.15
|
|
results: dict[str, str] = {}
|
|
|
|
def submit(char):
|
|
r = self.client.post("/api/character-generate-consistent-images", json=_payload(char))
|
|
results[char["name"]] = r.json()["job_id"]
|
|
|
|
threads = [threading.Thread(target=submit, args=(c,)) for c in (DWARF, ELF)]
|
|
[t.start() for t in threads]; [t.join() for t in threads]
|
|
self.assertNotEqual(results["Borin"], results["Lirael"])
|
|
self._assert_separate(self._wait(results["Borin"]), self._wait(results["Lirael"]))
|
|
|
|
def test_source_sheet_and_payload_are_not_mutated(self):
|
|
payload = _payload(DWARF)
|
|
before = repr(payload)
|
|
self.client.post("/api/character-generate-consistent-images", json=payload)
|
|
self.assertEqual(repr(payload), before)
|
|
self.assertEqual(repr(DWARF["sheet"]), repr(_payload(DWARF)["sheet"]))
|
|
|
|
def test_idempotency_key_never_returns_another_characters_job(self):
|
|
headers = {"Idempotency-Key": "reused-key"}
|
|
first = self.client.post("/api/character-generate-consistent-images", json=_payload(DWARF), headers=headers)
|
|
job_id = first.json()["job_id"]
|
|
same = self.client.post("/api/character-generate-consistent-images", json=_payload(DWARF), headers=headers)
|
|
self.assertEqual(same.json(), {"job_id": job_id, "deduplicated": True})
|
|
other = self.client.post("/api/character-generate-consistent-images", json=_payload(ELF), headers=headers)
|
|
self.assertEqual(other.status_code, 409)
|
|
self._wait(job_id)
|
|
self.assertEqual({c["style_key"].split("|")[3] for c in self.sheet_calls}, {"Borin"})
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|