From 4778373e328c6b4e4a3603a94a464cd6743ea219 Mon Sep 17 00:00:00 2001 From: akagifreeez Date: Tue, 7 Jul 2026 10:38:37 +0900 Subject: [PATCH 1/2] feat: add is_admin and is_moderator checks Adds two check decorators alongside the existing cooldown() in matrix/checks.py, following the same pattern: an inner async predicate that inspects ctx.room.power_levels (via nio's MatrixRoom.power_levels / PowerLevels.get_user_level) and a wrapper that registers it with Command.check() and returns the command. - is_admin(): requires power level >= 100 - is_moderator(): requires power level >= 50 Both feed into the existing Command._invoke -> CheckError flow, no new exception types needed. Closes #103, #104 Co-Authored-By: Claude --- matrix/__init__.py | 4 +- matrix/checks.py | 57 ++++++++++++++++ tests/test_checks.py | 155 +++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 215 insertions(+), 1 deletion(-) create mode 100644 tests/test_checks.py diff --git a/matrix/__init__.py b/matrix/__init__.py index 1845a21..0ebc2fe 100644 --- a/matrix/__init__.py +++ b/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", diff --git a/matrix/checks.py b/matrix/checks.py index d673ca7..ab0448a 100644 --- a/matrix/checks.py +++ b/matrix/checks.py @@ -2,6 +2,7 @@ if TYPE_CHECKING: from .command import Command + from .context import Context def cooldown(rate: int, period: float) -> Callable: @@ -27,3 +28,59 @@ def wrapper(cmd: "Command") -> "Command": return cmd return wrapper + + +def is_admin() -> Callable: + """ + Decorator to restrict a command to room admins (power level >= 100). + + ## 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: + return ctx.room.power_levels.get_user_level(ctx.sender) >= 100 # 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 >= 50). + + ## 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: + return ctx.room.power_levels.get_user_level(ctx.sender) >= 50 # type: ignore[no-any-return] + + def wrapper(cmd: "Command") -> "Command": + cmd.check(_is_moderator) + return cmd + + return wrapper diff --git a/tests/test_checks.py b/tests/test_checks.py new file mode 100644 index 0000000..5d3efc6 --- /dev/null +++ b/tests/test_checks.py @@ -0,0 +1,155 @@ +import pytest + +from unittest.mock import AsyncMock, Mock +from nio import MatrixRoom, RoomMessageText, AsyncClient, PowerLevels + +from matrix.checks import 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", + [ + (100, True), + (99, False), + (50, 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", + [ + (100, True), + (50, True), + (49, 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 + + +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", 50) + await cmd(ctx) + + assert called is True From 546c47c4ca04c7342df209f688172f65b6819b47 Mon Sep 17 00:00:00 2001 From: akagifreeez Date: Tue, 7 Jul 2026 10:41:31 +0900 Subject: [PATCH 2/2] test: cover users_default fallback for senders absent from power_levels Co-Authored-By: Claude Fable 5 --- tests/test_checks.py | 21 +++++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/tests/test_checks.py b/tests/test_checks.py index 5d3efc6..95303ce 100644 --- a/tests/test_checks.py +++ b/tests/test_checks.py @@ -97,6 +97,27 @@ async def my_command(ctx): 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