feat: add is_admin and is_moderator checks

matrix/__init__.py

@@ -13,7 +13,7 @@
 from .context import Context
 from .command import Command
 from .help import HelpCommand
-from .checks import cooldown
+from .checks import cooldown, is_admin, is_moderator
 from .room import Room
 from .space import Space
 from .message import Message
@@ -28,6 +28,8 @@
     "Context",
     "HelpCommand",
     "cooldown",
+    "is_admin",
+    "is_moderator",
     "Room",
     "Space",
     "Message",

matrix/checks.py

@@ -2,6 +2,11 @@
 
 if TYPE_CHECKING:
     from .command import Command
+    from .context import Context
+
+
+ADMIN_POWER_LEVEL: int = 100
+MODERATOR_POWER_LEVEL: int = 50
 
 
 def cooldown(rate: int, period: float) -> Callable:
@@ -27,3 +32,63 @@ def wrapper(cmd: "Command") -> "Command":
         return cmd
 
     return wrapper
+
+
+def is_admin() -> Callable:
+    """
+    Decorator to restrict a command to room admins
+    (power level >= `ADMIN_POWER_LEVEL`).
+
+    ## Example
+
+    ```python
+    @is_admin()
+    @bot.command("ban")
+    async def ban(ctx: Context, user_id: str) -> None:
+        await ctx.room.ban_user(user_id)
+
+    @ban.error(CheckError)
+    async def ban_error(ctx: Context, error: CheckError) -> None:
+        await ctx.reply("You must be an admin to use this command.")
+    ```
+    """
+
+    async def _is_admin(ctx: "Context") -> bool:
+        level = ctx.room.power_levels.get_user_level(ctx.sender)
+        return level >= ADMIN_POWER_LEVEL  # type: ignore[no-any-return]
+
+    def wrapper(cmd: "Command") -> "Command":
+        cmd.check(_is_admin)
+        return cmd
+
+    return wrapper
+
+
+def is_moderator() -> Callable:
+    """
+    Decorator to restrict a command to room moderators
+    (power level >= `MODERATOR_POWER_LEVEL`).
+
+    ## Example
+
+    ```python
+    @is_moderator()
+    @bot.command("kick")
+    async def kick(ctx: Context, user_id: str) -> None:
+        await ctx.room.kick_user(user_id)
+
+    @kick.error(CheckError)
+    async def kick_error(ctx: Context, error: CheckError) -> None:
+        await ctx.reply("You must be a moderator to use this command.")
+    ```
+    """
+
+    async def _is_moderator(ctx: "Context") -> bool:
+        level = ctx.room.power_levels.get_user_level(ctx.sender)
+        return level >= MODERATOR_POWER_LEVEL  # type: ignore[no-any-return]
+
+    def wrapper(cmd: "Command") -> "Command":
+        cmd.check(_is_moderator)
+        return cmd
+
+    return wrapper

tests/test_checks.py

@@ -0,0 +1,187 @@
+import pytest
+
+from unittest.mock import AsyncMock, Mock
+from nio import MatrixRoom, RoomMessageText, AsyncClient, PowerLevels
+
+from matrix.checks import (
+    ADMIN_POWER_LEVEL,
+    MODERATOR_POWER_LEVEL,
+    is_admin,
+    is_moderator,
+)
+from matrix.command import Command
+from matrix.context import Context
+from matrix.errors import CheckError
+from matrix.room import Room
+
+
+@pytest.fixture
+def client():
+    return AsyncMock(spec=AsyncClient)
+
+
+@pytest.fixture
+def bot(client):
+    bot = Mock()
+    bot.prefix = "!"
+    bot.client = client
+    bot.log = Mock()
+    bot.log.getChild = Mock(return_value=Mock())
+    return bot
+
+
+def make_event(sender: str) -> RoomMessageText:
+    return RoomMessageText.from_dict(
+        {
+            "content": {"body": "!restricted", "msgtype": "m.text"},
+            "event_id": "$event123",
+            "origin_server_ts": 123456,
+            "sender": sender,
+            "type": "m.room.message",
+        }
+    )
+
+
+def make_context(bot, client, sender: str, level: int) -> Context:
+    """Builds a real Context whose room reports `level` as the sender's power level."""
+    matrix_room = Mock(spec=MatrixRoom)
+    matrix_room.room_id = "!room:example.com"
+    matrix_room.power_levels = PowerLevels(users={sender: level})
+
+    room = Room(matrix_room, client)
+    return Context(bot, room, make_event(sender))
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "level, expected",
+    [
+        (ADMIN_POWER_LEVEL, True),
+        (ADMIN_POWER_LEVEL - 1, False),
+        (MODERATOR_POWER_LEVEL, False),
+        (0, False),
+    ],
+)
+async def test_is_admin_check__respects_power_level_boundaries(
+    bot, client, level, expected
+):
+    async def my_command(ctx):
+        pass
+
+    cmd = Command(my_command)
+    is_admin()(cmd)
+
+    check = cmd.checks[-1]
+    ctx = make_context(bot, client, "@user:example.com", level)
+
+    assert await check(ctx) is expected
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "level, expected",
+    [
+        (ADMIN_POWER_LEVEL, True),
+        (MODERATOR_POWER_LEVEL, True),
+        (MODERATOR_POWER_LEVEL - 1, False),
+        (0, False),
+    ],
+)
+async def test_is_moderator_check__respects_power_level_boundaries(
+    bot, client, level, expected
+):
+    async def my_command(ctx):
+        pass
+
+    cmd = Command(my_command)
+    is_moderator()(cmd)
+
+    check = cmd.checks[-1]
+    ctx = make_context(bot, client, "@user:example.com", level)
+
+    assert await check(ctx) is expected
+
+
+@pytest.mark.asyncio
+async def test_checks__fall_back_to_default_level_for_unlisted_sender(bot, client):
+    """A sender absent from power_levels.users gets users_default (0)."""
+    matrix_room = Mock(spec=MatrixRoom)
+    matrix_room.room_id = "!room:example.com"
+    matrix_room.power_levels = PowerLevels()
+
+    room = Room(matrix_room, client)
+    ctx = Context(bot, room, make_event("@user:example.com"))
+
+    for factory in (is_admin, is_moderator):
+
+        async def my_command(ctx):
+            pass
+
+        cmd = Command(my_command)
+        factory()(cmd)
+
+        assert await cmd.checks[-1](ctx) is False
+
+
+def test_is_admin__returns_the_same_command():
+    async def my_command(ctx):
+        pass
+
+    cmd = Command(my_command)
+    assert is_admin()(cmd) is cmd
+
+
+def test_is_moderator__returns_the_same_command():
+    async def my_command(ctx):
+        pass
+
+    cmd = Command(my_command)
+    assert is_moderator()(cmd) is cmd
+
+
+@pytest.mark.asyncio
+async def test_is_admin__raises_check_error_when_not_admin(bot, client):
+    called = False
+
+    async def restricted(ctx):
+        nonlocal called
+        called = True
+
+    cmd = Command(restricted)
+    is_admin()(cmd)
+
+    caught: list[Exception] = []
+
+    @cmd.error(CheckError)
+    async def on_check_error(ctx, error):
+        caught.append(error)
+
+    ctx = make_context(bot, client, "@user:example.com", 0)
+    await cmd(ctx)
+
+    assert called is False
+    assert len(caught) == 1
+    assert isinstance(caught[0], CheckError)
+
+
+@pytest.mark.asyncio
+async def test_is_moderator__allows_command_when_moderator(bot, client):
+    called = False
+
+    async def restricted(ctx):
+        nonlocal called
+        called = True
+
+    cmd = Command(restricted)
+    is_moderator()(cmd)
+
+    ctx = make_context(bot, client, "@mod:example.com", MODERATOR_POWER_LEVEL)
+    await cmd(ctx)
+
+    assert called is True
+
+
+def test_power_level_constants__match_matrix_defaults():
+    """The thresholds follow the defaults of `m.room.power_levels`."""
+    assert ADMIN_POWER_LEVEL == 100
+    assert MODERATOR_POWER_LEVEL == 50