136 lines
6.5 KiB
Python
136 lines
6.5 KiB
Python
"""Focused attribution parser regressions."""
|
||
from __future__ import annotations
|
||
|
||
import os
|
||
import json
|
||
import tempfile
|
||
import unittest
|
||
|
||
os.environ.setdefault("CONFIG_DIR", tempfile.mkdtemp(prefix="ttvc-conversation-test-config-"))
|
||
|
||
from routes.conversation import (
|
||
_attribution_parse,
|
||
_character_form_sheet,
|
||
_character_identity_fingerprint,
|
||
_cij_check_cancelled,
|
||
_cij_new,
|
||
_consistent_image_cancel_events,
|
||
_consistent_image_jobs,
|
||
_fallback_attribute_response,
|
||
_bundled_comfyui_workflow,
|
||
_comfyui_run,
|
||
)
|
||
|
||
|
||
class AttributionParserTests(unittest.TestCase):
|
||
def test_active_character_form_overlays_visual_identity_without_changing_record(self):
|
||
sheet = {
|
||
"physical": "base form",
|
||
"image_prompt": "base prompt",
|
||
"forms": [{"id": "wolf", "name": "Wolf form", "physical": "black wolf", "image_prompt": "black wolf portrait"}],
|
||
"active_form_id": "wolf",
|
||
}
|
||
resolved = _character_form_sheet(sheet)
|
||
self.assertEqual(resolved["physical"], "black wolf")
|
||
self.assertEqual(resolved["image_prompt"], "black wolf portrait")
|
||
self.assertEqual(sheet["active_form_id"], "wolf")
|
||
self.assertNotEqual(_character_identity_fingerprint(sheet), _character_identity_fingerprint({**sheet, "active_form_id": ""}))
|
||
|
||
def test_fallback_accepts_mixed_guillemets(self):
|
||
result = _fallback_attribute_response("Intro. »Ein langer Satz endet hier.‹ Danach.")
|
||
self.assertEqual([segment["type"] for segment in result["segments"]], ["narration", "dialogue", "narration"])
|
||
self.assertIn("Ein langer Satz endet hier.", result["segments"][1]["text"])
|
||
|
||
def test_repairs_epigraph_split_before_mixed_guillemet_close(self):
|
||
raw = '{"segments":[' \
|
||
'{"speaker":"Schwester Larissa","type":"dialogue","text":"»Von dem Windhager Felsgesteyn bis zur Kueste von Mendena ziehet es sich hin, von Ifirns Eisgebirg bis zu dem dampfenden Waldt der Mohas, und"},' \
|
||
'{"speaker":"Narrator","type":"narration","text":"fürwahr, es ist ein seltsam Land. Lieblich bisweilen, mild und voll der guten Menschen. ‹ Zitiert aus dem Buch."}' \
|
||
']}'
|
||
result = _attribution_parse(raw, "unused")
|
||
self.assertEqual(len(result["segments"]), 2)
|
||
dialogue, narration = result["segments"]
|
||
self.assertEqual(dialogue["speaker"], "Schwester Larissa")
|
||
self.assertEqual(dialogue["type"], "dialogue")
|
||
self.assertTrue(dialogue["text"].endswith("Land. Lieblich bisweilen, mild und voll der guten Menschen. ‹"))
|
||
self.assertEqual(narration["text"], "Zitiert aus dem Buch.")
|
||
|
||
def test_book2_neutral_fixture_preserves_multiple_speakers(self):
|
||
fixture_path = os.path.join(os.path.dirname(__file__), "fixtures", "book2_attribution.json")
|
||
with open(fixture_path, encoding="utf-8") as handle:
|
||
fixture = json.load(handle)
|
||
result = _attribution_parse(fixture["raw"], " ".join(fixture["known_characters"]))
|
||
self.assertEqual(
|
||
[segment["speaker"] for segment in result["segments"]],
|
||
["Narrator", "Captain Vale", "The Ferryman", "Mira"],
|
||
)
|
||
self.assertTrue(all(segment["type"] == "dialogue" for segment in result["segments"][1:]))
|
||
|
||
def test_bundled_comfyui_workflows_have_required_prompt_and_save_nodes(self):
|
||
for image_kind in ("portrait", "concept_art"):
|
||
raw, prompt_node, prompt_field, save_node = _bundled_comfyui_workflow(image_kind)
|
||
graph = json.loads(raw)
|
||
self.assertIn(prompt_node, graph)
|
||
self.assertEqual(graph[prompt_node]["class_type"], "CLIPTextEncode")
|
||
self.assertIn(prompt_field, graph[prompt_node]["inputs"])
|
||
self.assertIn(save_node, graph)
|
||
self.assertIn("filename_prefix", graph[save_node]["inputs"])
|
||
|
||
def test_comfyui_submission_injects_prompt_polls_and_returns_image(self):
|
||
from unittest.mock import patch
|
||
|
||
class ResponseStub:
|
||
def __init__(self, payload=None, content=b"", headers=None):
|
||
self._payload, self.content, self.headers = payload, content, headers or {}
|
||
self.ok = True
|
||
self.status_code = 200
|
||
self.text = ""
|
||
def json(self):
|
||
return self._payload
|
||
def raise_for_status(self):
|
||
return None
|
||
|
||
graph = {
|
||
"2": {"class_type": "CLIPTextEncode", "inputs": {"text": "old"}},
|
||
"5": {"class_type": "KSampler", "inputs": {"seed": 0}},
|
||
"7": {"class_type": "SaveImage", "inputs": {"filename_prefix": "test"}},
|
||
}
|
||
post = ResponseStub({"prompt_id": "p1"})
|
||
history = ResponseStub({"p1": {"status": {"completed": True}, "outputs": {"7": {"images": [{"filename": "out.png"}]}}}})
|
||
view = ResponseStub(content=b"png", headers={"Content-Type": "image/png"})
|
||
with patch("routes.conversation.requests.post", return_value=post) as post_call, \
|
||
patch("routes.conversation.requests.get", side_effect=[history, view]), \
|
||
patch("routes.conversation.time.sleep"):
|
||
result = _comfyui_run("http://comfyui:8188", json.dumps(graph), "2", "text", "7", "new prompt", "stable")
|
||
self.assertTrue(result.startswith("data:image/png;base64,"))
|
||
submitted = post_call.call_args.kwargs["json"]["prompt"]
|
||
self.assertEqual(submitted["2"]["inputs"]["text"], "new prompt")
|
||
self.assertNotEqual(submitted["5"]["inputs"]["seed"], 0)
|
||
|
||
def test_consistent_image_job_can_be_cancelled(self):
|
||
from fastapi.testclient import TestClient
|
||
from server import app
|
||
|
||
job_id = "cancel-test"
|
||
_cij_new(job_id, "Test character")
|
||
try:
|
||
with TestClient(app) as client:
|
||
response = client.post(
|
||
f"/api/character-generate-consistent-images/cancel/{job_id}",
|
||
headers={"origin": "http://testserver"},
|
||
)
|
||
self.assertEqual(response.status_code, 200)
|
||
self.assertTrue(response.json()["cancelled"])
|
||
with self.assertRaises(Exception):
|
||
_cij_check_cancelled(job_id)
|
||
finally:
|
||
_consistent_image_jobs.pop(job_id, None)
|
||
_consistent_image_cancel_events.pop(job_id, None)
|
||
from core.database import _open
|
||
with _open() as conn:
|
||
conn.execute("DELETE FROM background_jobs WHERE job_id=?", (job_id,))
|
||
conn.commit()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
unittest.main()
|