Files
technical-visualization-haness/tests/test_prompt.py
T

42 lines
1.2 KiB
Python

from __future__ import annotations
import json
import unittest
from pathlib import Path
from techviz.document import build_context
from techviz.prompt import build_agent_prompt
from techviz.spec import validate_raw_spec
ROOT = Path(__file__).resolve().parents[1]
class PromptTests(unittest.TestCase):
def test_embedded_vizspec_scaffold_is_valid_and_context_bound(self) -> None:
context = build_context(
ROOT / "examples/docs/payment-flow.md",
marker_id="payment-request",
)
prompt = build_agent_prompt(context)
section_start = prompt.index("## VizSpec 1.0 shape")
json_start = prompt.index("{\n", section_start)
json_end = prompt.index("\n\nFor a sequence diagram", json_start)
scaffold = json.loads(prompt[json_start:json_end])
validate_raw_spec(scaffold)
self.assertEqual(
scaffold["source_context"],
{
"document": context["document"],
"document_sha256": context["document_sha256"],
"anchor": context["anchor"],
},
)
self.assertFalse(scaffold["groups"])
self.assertFalse(scaffold["legend"])
if __name__ == "__main__":
unittest.main()