nonebot2/tests/test_adapters/test_template.py
Johnny Hsieh b65b3b438c
🐛 Fix: MessageTemplate 禁止访问私有属性 (#2509)
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2024-01-04 11:11:37 +08:00

77 lines
2.4 KiB
Python

import pytest
from nonebot.adapters import MessageTemplate
from utils import FakeMessage, FakeMessageSegment, escape_text
def test_template_basis():
template = MessageTemplate("{key:.3%}")
formatted = template.format(key=0.123456789)
assert formatted == "12.346%"
def test_template_message():
template = FakeMessage.template("{a:custom}{b:text}{c:image}/{d}")
@template.add_format_spec
def custom(input: str) -> str:
return f"{input}-custom!"
with pytest.raises(ValueError, match="already exists"):
template.add_format_spec(custom)
format_args = {
"a": "custom",
"b": "text",
"c": "https://example.com/test",
"d": 114,
}
formatted = template.format(**format_args)
assert template.format_map(format_args) == formatted
assert formatted.extract_plain_text() == "custom-custom!text/114"
assert str(formatted) == "custom-custom!text[fake:image]/114"
def test_rich_template_message():
pic1, pic2, pic3 = (
FakeMessageSegment.image("file:///pic1.jpg"),
FakeMessageSegment.image("file:///pic2.jpg"),
FakeMessageSegment.image("file:///pic3.jpg"),
)
template = FakeMessage.template("{}{}" + pic2 + "{}")
result = template.format(pic1, "[fake:image]", pic3)
assert result["image"] == FakeMessage([pic1, pic2, pic3])
assert str(result) == (
"[fake:image]" + escape_text("[fake:image]") + "[fake:image]" + "[fake:image]"
)
def test_message_injection():
template = FakeMessage.template("{name}Is Bad")
message = template.format(name="[fake:image]")
assert message.extract_plain_text() == escape_text("[fake:image]Is Bad")
def test_malformed_template():
positive_template = FakeMessage.template("{a}{b}")
message = positive_template.format(a="a", b="b")
assert message.extract_plain_text() == "ab"
malformed_template = FakeMessage.template("{a.__init__}")
with pytest.raises(ValueError, match="private attribute"):
message = malformed_template.format(a="a")
malformed_template = FakeMessage.template("{a[__builtins__]}")
with pytest.raises(ValueError, match="private attribute"):
message = malformed_template.format(a=globals())
malformed_template = MessageTemplate(
"{a[__builtins__][__import__]}{b.__init__}", private_getattr=True
)
message = malformed_template.format(a=globals(), b="b")