nonebot2/tests/test_adapters/test_template.py

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")