tests/agents/d3/test_registry.py

from pydantic import BaseModel
from app.agents.d3.registry import D3Template, register, TEMPLATES, catalog_for_prompt


class _Sample(BaseModel):
    value: int


class _StrSample(BaseModel):
    label: str


def _make():
    return D3Template(
        id="dummy", family="chart", title="Dummy",
        when_to_use="never", data_requirements="an int",
        schema=_Sample, html_template="<b>__DATA__</b>",
        golden_sample=_Sample(value=7),
    )


def test_render_injects_data():
    t = _make()
    assert t.render(_Sample(value=7)) == '<b>{"value":7}</b>'


def test_register_and_catalog():
    register(_make())
    assert "dummy" in TEMPLATES
    line = catalog_for_prompt()
    assert "dummy" in line and "chart" in line and "an int" in line


def test_render_escapes_script_close_tag_in_data():
    t = D3Template(
        id="dummy_script", family="chart", title="Dummy",
        when_to_use="never", data_requirements="a string",
        schema=_StrSample, html_template="<script>var d = __DATA__;</script>",
        golden_sample=_StrSample(label="x"),
    )
    malicious = _StrSample(label="</script><script>alert(1)</script>")
    rendered = t.render(malicious)

    # Only the template's own two real <script> tags should remain;
    # the injected data must not contribute an unescaped </script>.
    assert rendered.count("</script>") == 1
    assert '<\\/script>' in rendered