diff --git a/matrix/message.py b/matrix/message.py index 0e90d6e..ed1aaee 100644 --- a/matrix/message.py +++ b/matrix/message.py @@ -1,6 +1,11 @@ from typing import TYPE_CHECKING, Self -from nio import AsyncClient, Event +from nio import ( + AsyncClient, + Event, + RoomGetStateEventError, + RoomGetStateEventResponse, +) from matrix.types import Reaction from matrix.content import ReactionContent, EditContent @@ -175,3 +180,70 @@ async def oops(ctx: Context): ), error_message="Failed to delete message", ) + + async def _fetch_pinned(self) -> list[str]: + response = await self.client.room_get_state_event( + room_id=self.room.room_id, + event_type="m.room.pinned_events", + ) + + if isinstance(response, RoomGetStateEventResponse): + return list(response.content.get("pinned", [])) + if isinstance(response, RoomGetStateEventError): + if response.status_code == "404": + return [] + raise MatrixError(f"Failed to fetch pinned events: {response.message}") + + raise MatrixError("Unexpected response type when fetching pinned events") + + async def pin(self) -> None: + """Pin this message to the room. + + ## Example + ```python + @bot.command() + async def pin(ctx: Context): + await ctx.message.pin() + ``` + """ + pinned = await self._fetch_pinned() + + if self.event_id in pinned: + return + + pinned.append(self.event_id) + + await matrix_call( + self.client.room_put_state( + room_id=self.room.room_id, + event_type="m.room.pinned_events", + content={"pinned": pinned}, + ), + error_message="Failed to pin message", + ) + + async def unpin(self) -> None: + """Unpin this message from the room. + + ## Example + ```python + @bot.command() + async def unpin(ctx: Context): + await ctx.message.unpin() + ``` + """ + pinned = await self._fetch_pinned() + + if self.event_id not in pinned: + return + + pinned.remove(self.event_id) + + await matrix_call( + self.client.room_put_state( + room_id=self.room.room_id, + event_type="m.room.pinned_events", + content={"pinned": pinned}, + ), + error_message="Failed to unpin message", + ) diff --git a/tests/test_message.py b/tests/test_message.py index b7a4b13..2bbcbdb 100644 --- a/tests/test_message.py +++ b/tests/test_message.py @@ -4,6 +4,9 @@ MatrixRoom, AsyncClient, Event, + RoomGetStateEventResponse, + RoomGetStateEventError, + RoomPutStateError, RoomSendError, RoomRedactError, ) @@ -257,3 +260,138 @@ def test_message_client__expect_correct_reference(message, client): def test_message_repr__expect_formatted_string(message): result = repr(message) assert result == "" + + +@pytest.mark.asyncio +async def test_pin_no_existing_pinned_events__expect_state_updated(message, client): + mock_get_error = MagicMock(spec=RoomGetStateEventError) + mock_get_error.status_code = "404" + client.room_get_state_event = AsyncMock(return_value=mock_get_error) + + client.room_put_state = AsyncMock(return_value=MagicMock()) + + await message.pin() + + client.room_get_state_event.assert_awaited_once_with( + room_id="!room:example.com", + event_type="m.room.pinned_events", + ) + client.room_put_state.assert_awaited_once_with( + room_id="!room:example.com", + event_type="m.room.pinned_events", + content={"pinned": ["$event123"]}, + ) + + +@pytest.mark.asyncio +async def test_pin_with_existing_pinned_events__expect_appended_state_updated( + message, client +): + mock_get_resp = MagicMock(spec=RoomGetStateEventResponse) + mock_get_resp.content = {"pinned": ["$other_event"]} + client.room_get_state_event = AsyncMock(return_value=mock_get_resp) + + client.room_put_state = AsyncMock(return_value=MagicMock()) + + await message.pin() + + client.room_put_state.assert_awaited_once_with( + room_id="!room:example.com", + event_type="m.room.pinned_events", + content={"pinned": ["$other_event", "$event123"]}, + ) + + +@pytest.mark.asyncio +async def test_pin_already_pinned__expect_no_state_update(message, client): + mock_get_resp = MagicMock(spec=RoomGetStateEventResponse) + mock_get_resp.content = {"pinned": ["$event123"]} + client.room_get_state_event = AsyncMock(return_value=mock_get_resp) + + client.room_put_state = AsyncMock() + + await message.pin() + + client.room_put_state.assert_not_called() + + +@pytest.mark.asyncio +async def test_pin_with_forbidden_error__expect_matrix_error(message, client): + mock_get_error = MagicMock(spec=RoomGetStateEventError) + mock_get_error.status_code = "403" + mock_get_error.message = "Forbidden" + client.room_get_state_event = AsyncMock(return_value=mock_get_error) + + with pytest.raises(MatrixError, match="Failed to fetch pinned events: Forbidden"): + await message.pin() + + +@pytest.mark.asyncio +async def test_pin_with_put_error__expect_matrix_error(message, client): + mock_get_error = MagicMock(spec=RoomGetStateEventError) + mock_get_error.status_code = "404" + client.room_get_state_event = AsyncMock(return_value=mock_get_error) + + client.room_put_state = AsyncMock( + return_value=RoomPutStateError("not allowed", "M_FORBIDDEN") + ) + + with pytest.raises(MatrixError, match="Failed to pin message: .*M_FORBIDDEN"): + await message.pin() + + +@pytest.mark.asyncio +async def test_unpin_event_is_pinned__expect_state_updated(message, client): + mock_get_resp = MagicMock(spec=RoomGetStateEventResponse) + mock_get_resp.content = {"pinned": ["$other_event", "$event123"]} + client.room_get_state_event = AsyncMock(return_value=mock_get_resp) + + client.room_put_state = AsyncMock(return_value=MagicMock()) + + await message.unpin() + + client.room_put_state.assert_awaited_once_with( + room_id="!room:example.com", + event_type="m.room.pinned_events", + content={"pinned": ["$other_event"]}, + ) + + +@pytest.mark.asyncio +async def test_unpin_event_not_pinned__expect_no_state_update(message, client): + mock_get_resp = MagicMock(spec=RoomGetStateEventResponse) + mock_get_resp.content = {"pinned": ["$other_event"]} + client.room_get_state_event = AsyncMock(return_value=mock_get_resp) + + client.room_put_state = AsyncMock() + + await message.unpin() + + client.room_put_state.assert_not_called() + + +@pytest.mark.asyncio +async def test_unpin_no_pinned_events__expect_no_state_update(message, client): + mock_get_error = MagicMock(spec=RoomGetStateEventError) + mock_get_error.status_code = "404" + client.room_get_state_event = AsyncMock(return_value=mock_get_error) + + client.room_put_state = AsyncMock() + + await message.unpin() + + client.room_put_state.assert_not_called() + + +@pytest.mark.asyncio +async def test_unpin_with_put_error__expect_matrix_error(message, client): + mock_get_resp = MagicMock(spec=RoomGetStateEventResponse) + mock_get_resp.content = {"pinned": ["$event123"]} + client.room_get_state_event = AsyncMock(return_value=mock_get_resp) + + client.room_put_state = AsyncMock( + return_value=RoomPutStateError("not allowed", "M_FORBIDDEN") + ) + + with pytest.raises(MatrixError, match="Failed to unpin message: .*M_FORBIDDEN"): + await message.unpin()