Spaces:
Sleeping
Sleeping
File size: 1,402 Bytes
2e818da | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 | 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
|