disnake

1
# SPDX-License-Identifier: MIT
2

3
import contextlib
4
from typing import Any, Iterator, Type, TypeVar
5

6
import pytest
7
from typing_extensions import assert_type
8

9
from disnake import ui
10
from disnake.ui.button import V_co
11

12
T = TypeVar("T", bound=ui.Item)
13

14

15
@contextlib.contextmanager
16
def create_callback(item_type: Type[T]) -> Iterator["ui.item.ItemCallbackType[T]"]:
17
    async def callback(self, item, inter) -> None:
18
        pytest.fail("callback should not be invoked")
19

20
    yield callback
21

22
    # ensure instantiation works
23
    assert callback.__discord_ui_model_type__(**callback.__discord_ui_model_kwargs__)
24

25

26
class _CustomButton(ui.Button[V_co]):
27
    def __init__(self, *, param: float = 42.0) -> None:
28
        pass
29

30

31
class TestDecorator:
32
    def test_default(self) -> None:
33
        with create_callback(ui.Button[ui.View]) as func:
34
            res = ui.button(custom_id="123")(func)
35
            assert_type(res, ui.item.DecoratedItem[ui.Button[ui.View]])
36

37
            assert func.__discord_ui_model_type__ is ui.Button
38
            assert func.__discord_ui_model_kwargs__ == {"custom_id": "123"}
39

40
        with create_callback(ui.StringSelect[ui.View]) as func:
41
            res = ui.string_select(custom_id="123")(func)
42
            assert_type(res, ui.item.DecoratedItem[ui.StringSelect[ui.View]])
43

44
            assert func.__discord_ui_model_type__ is ui.StringSelect
45
            assert func.__discord_ui_model_kwargs__ == {"custom_id": "123"}
46

47
    # from here on out we're mostly only testing the button decorator,
48
    # as @ui.string_select etc. works identically
49

50
    @pytest.mark.parametrize("cls", [_CustomButton, _CustomButton[Any]])
51
    def test_cls(self, cls: Type[_CustomButton]) -> None:
52
        with create_callback(cls) as func:
53
            res = ui.button(cls=cls, param=1337)(func)
54
            assert_type(res, ui.item.DecoratedItem[cls])
55

56
            # should strip to origin type
57
            assert func.__discord_ui_model_type__ is _CustomButton
58
            assert func.__discord_ui_model_kwargs__ == {"param": 1337}
59

60
    # typing-only check
61
    def _test_typing_cls(self) -> None:
62
        ui.button(
63
            cls=_CustomButton,
64
            this_should_not_work="h",  # type: ignore
65
        )
66

67
    @pytest.mark.parametrize(
68
        ("decorator", "invalid_cls"),
69
        [
70
            (ui.button, ui.StringSelect),
71
            (ui.string_select, ui.Button),
72
            (ui.user_select, ui.Button),
73
            (ui.role_select, ui.Button),
74
            (ui.mentionable_select, ui.Button),
75
            (ui.channel_select, ui.Button),
76
        ],
77
    )
78
    def test_cls_invalid(self, decorator, invalid_cls) -> None:
79
        for cls in [123, int, invalid_cls]:
80
            with pytest.raises(TypeError, match=r"cls argument must be"):
81
                decorator(cls=cls)
82