From 86aa44b2603db7089207b910e37345099b0116b5 Mon Sep 17 00:00:00 2001 From: HgO Date: Mon, 22 Jan 2024 11:35:13 +0100 Subject: [PATCH] implement matrix client failover --- .gitignore | 2 +- config.sample.yaml | 37 ++-- matrix_alertbot/alertmanager.py | 45 ++-- matrix_alertbot/callback.py | 211 ++++++++++++------ matrix_alertbot/chat_functions.py | 15 +- matrix_alertbot/command.py | 25 ++- matrix_alertbot/config.py | 65 ++++-- matrix_alertbot/errors.py | 6 + matrix_alertbot/main.py | 167 ++------------ matrix_alertbot/matrix.py | 252 ++++++++++++++++++++++ matrix_alertbot/webhook.py | 36 +++- setup.cfg | 4 +- tests/resources/config/config.full.yml | 59 +++-- tests/resources/config/config.minimal.yml | 13 +- tests/test_callback.py | 199 +++++------------ tests/test_chat_functions.py | 19 +- tests/test_command.py | 4 - tests/test_config.py | 54 +++-- tests/test_webhook.py | 8 +- tests/utils.py | 22 -- 20 files changed, 722 insertions(+), 521 deletions(-) create mode 100644 matrix_alertbot/matrix.py delete mode 100644 tests/utils.py diff --git a/.gitignore b/.gitignore index d77149f..aa47cc9 100644 --- a/.gitignore +++ b/.gitignore @@ -12,7 +12,7 @@ env3/ *.db store/ cache/ -token.json +*token.json # Config file config.yaml diff --git a/config.sample.yaml b/config.sample.yaml index 68f3d4c..81c6c5e 100644 --- a/config.sample.yaml +++ b/config.sample.yaml @@ -7,28 +7,29 @@ command_prefix: "!alert" # Options for connecting to the bot's Matrix account matrix: - # The Matrix User ID of the bot account - user_id: "@bot:matrix.example.com" + accounts: + - # The Matrix User IDs of the bot account + id: "@bot:matrix.example.com" - # Matrix account password (optional if access token used) - user_password: "password" + # Matrix account password (optional if access token used) + password: "password" - # Matrix account access token (optional if password used) - # If not set, the server will provide an access token after log in, - # which will be stored in the user token file (see below) - #user_token: "" + # Matrix account access token (optional if password used) + # If not set, the server will provide an access token after log in, + # which will be stored in the user token file (see below) + #token: "" - # Path to the file where to store the user access token - user_token_file: "token.json" + # Path to the file where to store the user access token + token_file: "token.json" - # The URL of the homeserver to connect to - url: https://matrix.example.com + # The URL of the homeserver to connect to + url: https://matrix.example.com - # The device ID that is **non pre-existing** device - # If this device ID already exists, messages will be dropped silently in encrypted rooms - # If not set the server will provide a device ID after log in. Note that this ID - # will change each time the bot reconnects. - # device_id: ABCDEFGHIJ + # The device ID that is **non pre-existing** device + # If this device ID already exists, messages will be dropped silently in encrypted rooms + # If not set the server will provide a device ID after log in. Note that this ID + # will change each time the bot reconnects. + # device_id: ABCDEFGHIJ # What to name the logged in device device_name: matrix-alertbot @@ -79,6 +80,8 @@ logging: file_logging: # Whether logging to a file is enabled enabled: false + # Logging level specific to file logging (optional) + level: WARN # The path to the file to log to. May be relative or absolute filepath: matrix-alertbot.log # Configure logging to the console output diff --git a/matrix_alertbot/alertmanager.py b/matrix_alertbot/alertmanager.py index 905f61c..a2535e9 100644 --- a/matrix_alertbot/alertmanager.py +++ b/matrix_alertbot/alertmanager.py @@ -2,7 +2,7 @@ from __future__ import annotations import logging from datetime import datetime, timedelta -from typing import Dict, List, Optional +from typing import Any, Dict, List, Optional, Tuple, TypedDict, cast import aiohttp from aiohttp import ClientError @@ -23,6 +23,24 @@ MAX_DURATION = timedelta(days=3652) logger = logging.getLogger(__name__) +AlertDict = TypedDict( + "AlertDict", + { + "fingerprint": str, + "labels": Dict[str, str], + }, +) + +SilenceDict = TypedDict( + "SilenceDict", + { + "id": str, + "matchers": List[Dict[str, Any]], + "createdBy": str, + "status": Dict[str, str], + }, +) + class AlertmanagerClient: def __init__(self, url: str, cache: Cache) -> None: @@ -33,7 +51,7 @@ class AlertmanagerClient: async def close(self) -> None: await self.session.close() - async def get_alerts(self) -> List[Dict]: + async def get_alerts(self) -> List[AlertDict]: try: async with self.session.get(f"{self.api_url}/alerts") as response: response.raise_for_status() @@ -43,12 +61,12 @@ class AlertmanagerClient: "Cannot fetch alerts from Alertmanager" ) from e - async def get_alert(self, fingerprint: str) -> Dict: + async def get_alert(self, fingerprint: str) -> AlertDict: logger.debug(f"Fetching details for alert with fingerprint {fingerprint}") alerts = await self.get_alerts() return self._find_alert(fingerprint, alerts) - async def get_silences(self) -> List[Dict]: + async def get_silences(self) -> List[SilenceDict]: try: async with self.session.get(f"{self.api_url}/silences") as response: response.raise_for_status() @@ -58,7 +76,7 @@ class AlertmanagerClient: "Cannot fetch silences from Alertmanager" ) from e - async def get_silence(self, silence_id: str) -> Dict: + async def get_silence(self, silence_id: str) -> SilenceDict: logger.debug(f"Fetching details for silence with ID {silence_id}") silences = await self.get_silences() return self._find_silence(silence_id, silences) @@ -93,12 +111,15 @@ class AlertmanagerClient: logger.debug( f"Reading silence for alert with fingerprint {fingerprint} from cache" ) - try: - silence_id: Optional[str] - expire_time: Optional[int] - silence_id, expire_time = self.cache.get(fingerprint, expire_time=True) - except TypeError: + + cache_result = cast( + Optional[Tuple[str, int]], self.cache.get(fingerprint, expire_time=True) + ) + if cache_result is not None: + silence_id, expire_time = cache_result + else: silence_id = None + expire_time = None if silence_id is None: raise SilenceNotFoundError( @@ -202,14 +223,14 @@ class AlertmanagerClient: ) from e @staticmethod - def _find_alert(fingerprint: str, alerts: List[Dict]) -> Dict: + def _find_alert(fingerprint: str, alerts: List[AlertDict]) -> AlertDict: for alert in alerts: if alert["fingerprint"] == fingerprint: return alert raise AlertNotFoundError(f"Cannot find alert with fingerprint {fingerprint}") @staticmethod - def _find_silence(silence_id: str, silences: List[Dict]) -> Dict: + def _find_silence(silence_id: str, silences: List[SilenceDict]) -> SilenceDict: for silence in silences: if silence["id"] == silence_id: return silence diff --git a/matrix_alertbot/callback.py b/matrix_alertbot/callback.py index 2f57cb4..f2963e1 100644 --- a/matrix_alertbot/callback.py +++ b/matrix_alertbot/callback.py @@ -1,25 +1,27 @@ +from __future__ import annotations + import logging from diskcache import Cache -from nio import ( - AsyncClient, +from nio.client import AsyncClient +from nio.events import ( + Event, InviteMemberEvent, - JoinError, KeyVerificationCancel, KeyVerificationKey, KeyVerificationMac, KeyVerificationStart, - LocalProtocolError, - MatrixRoom, MegolmEvent, + ReactionEvent, RedactionEvent, - RoomGetEventError, RoomMessageText, - SendRetryError, - ToDeviceError, - UnknownEvent, + RoomMessageUnknown, ) +from nio.exceptions import LocalProtocolError, SendRetryError +from nio.responses import JoinError, RoomGetEventError, RoomSendError, ToDeviceError +from nio.rooms import MatrixRoom +import matrix_alertbot.matrix from matrix_alertbot.alertmanager import AlertmanagerClient from matrix_alertbot.chat_functions import strip_fallback from matrix_alertbot.command import AckAlertCommand, CommandFactory, UnackAlertCommand @@ -35,6 +37,7 @@ class Callbacks: alertmanager_client: AlertmanagerClient, cache: Cache, config: Config, + matrix_client_pool: matrix_alertbot.matrix.MatrixClientPool, ): """ Args: @@ -47,6 +50,7 @@ class Callbacks: config: Bot configuration parameters. """ self.matrix_client = matrix_client + self.matrix_client_pool = matrix_client_pool self.cache = cache self.alertmanager_client = alertmanager_client self.config = config @@ -60,8 +64,12 @@ class Callbacks: event: The event defining the message. """ + # Ignore message when we aren't the leader in the client pool + if self.matrix_client is not self.matrix_client_pool.matrix_client: + return + # Ignore messages from ourselves - if event.sender == self.matrix_client.user: + if event.sender in self.config.user_ids: return # Ignore messages from unauthorized room @@ -72,13 +80,16 @@ class Callbacks: msg = strip_fallback(event.body) logger.debug( - f"Bot message received for room {room.display_name} | " - f"{room.user_name(event.sender)}: {msg}" + f"Bot {self.matrix_client.user_id} | Room ID {room.room_id} | " + f"Event ID {event.event_id} | Sender {event.sender} | " + f"Message received: {msg}" ) # Process as message if in a public room without command prefix has_command_prefix = msg.startswith(self.command_prefix) if not has_command_prefix: logger.debug( + f"Bot {self.matrix_client.user_id} | Room ID {room.room_id} | " + f"Event ID {event.event_id} | Sender {event.sender} | " f"Cannot process message: Command prefix {self.command_prefix} not provided." ) return @@ -91,7 +102,11 @@ class Callbacks: ) if reacted_to_event_id is not None: - logger.debug(f"Command in reply to event ID {reacted_to_event_id}") + logger.debug( + f"Bot {self.matrix_client.user_id} | Room ID {room.room_id} | " + f"Event ID {event.event_id} | Sender {event.sender} | " + f"Command received is in reply to event ID {reacted_to_event_id}" + ) # Remove the command prefix cmd = msg[len(self.command_prefix) :] @@ -108,13 +123,22 @@ class Callbacks: reacted_to_event_id, ) except TypeError as e: - logging.error(f"Cannot process command '{cmd}': {e}") + logger.error( + f"Bot {self.matrix_client.user_id} | Room ID {room.room_id} | " + f"Event ID {event.event_id} | Sender {event.sender} | " + f"Cannot process command '{cmd}': {e}" + ) return try: await command.process() except (SendRetryError, LocalProtocolError) as e: - logger.exception(f"Unable to send message to {room.room_id}", exc_info=e) + logger.exception( + f"Bot {self.matrix_client.user_id} | Room ID {room.room_id} | " + f"Event ID {event.event_id} | Sender {event.sender} | " + f"Cannot send message to room.", + exc_info=e, + ) async def invite(self, room: MatrixRoom, event: InviteMemberEvent) -> None: """Callback for when an invite is received. Join the room specified in the invite. @@ -128,7 +152,11 @@ class Callbacks: if room.room_id not in self.config.allowed_rooms: return - logger.debug(f"Got invite to {room.room_id} from {event.sender}.") + logger.debug( + f"Bot {self.matrix_client.user_id} | Room ID {room.room_id} | " + f"Sender {event.sender} | " + f"Invitation received." + ) # Attempt to join 3 times before giving up for attempt in range(3): @@ -145,7 +173,11 @@ class Callbacks: logger.error("Unable to join room: %s", room.room_id) # Successfully joined room - logger.info(f"Joined {room.room_id}") + logger.info( + f"Bot {self.matrix_client.user_id} | Room ID {room.room_id} | " + f"Sender {event.sender} | " + f"Room joined." + ) async def invite_event_filtered_callback( self, room: MatrixRoom, event: InviteMemberEvent @@ -160,9 +192,7 @@ class Callbacks: # This is our own membership (invite) event await self.invite(room, event) - async def _reaction( - self, room: MatrixRoom, event: UnknownEvent, alert_event_id: str - ) -> None: + async def reaction(self, room: MatrixRoom, event: ReactionEvent) -> None: """A reaction was sent to one of our messages. Let's send a reply acknowledging it. Args: @@ -172,34 +202,48 @@ class Callbacks: reacted_to_id: The event ID that the reaction points to. """ + # Ignore message when we aren't the leader in the client pool + if self.matrix_client is not self.matrix_client_pool.matrix_client: + return + # Ignore reactions from unauthorized room if room.room_id not in self.config.allowed_rooms: return # Ignore reactions from ourselves - if event.sender == self.matrix_client.user: + if event.sender in self.config.user_ids: return - reaction = event.source.get("content", {}).get("m.relates_to", {}).get("key") - logger.debug(f"Got reaction {reaction} to {room.room_id} from {event.sender}.") + logger.debug( + f"Bot {self.matrix_client.user_id} | Room ID {room.room_id} | " + f"Event ID {event.event_id} | Sender {event.sender} | " + f"Reaction received: {event.key}" + ) - if reaction not in self.config.allowed_reactions: - logger.warning(f"Uknown duration reaction {reaction}") + if event.key not in self.config.allowed_reactions: + logger.warning( + f"Bot {self.matrix_client.user_id} | Room ID {room.room_id} | " + f"Event ID {event.event_id} | Sender {event.sender} | " + f"Reaction not handled: {event.key}" + ) return + alert_event_id = event.reacts_to # Get the original event that was reacted to event_response = await self.matrix_client.room_get_event( room.room_id, alert_event_id ) if isinstance(event_response, RoomGetEventError): logger.warning( - f"Error getting event that was reacted to ({alert_event_id})" + f"Bot {self.matrix_client.user_id} | Room ID {room.room_id} | " + f"Event ID {event.event_id} | Sender {event.sender} | " + f"Cannot get event related to the reaction and with event ID {alert_event_id}" ) return reacted_to_event = event_response.event # Only acknowledge reactions to events that we sent - if reacted_to_event.sender != self.config.user_id: + if reacted_to_event.sender not in self.config.user_ids: return # Send a message acknowledging the reaction @@ -217,18 +261,31 @@ class Callbacks: try: await command.process() except (SendRetryError, LocalProtocolError) as e: - logger.exception(f"Unable to send message to {room.room_id}", exc_info=e) + logger.exception( + f"Bot {self.matrix_client.user_id} | Room ID {room.room_id} | " + f"Event ID {event.event_id} | Sender {event.sender} | " + f"Cannot send message to room.", + exc_info=e, + ) async def redaction(self, room: MatrixRoom, event: RedactionEvent) -> None: + # Ignore message when we aren't the leader in the client pool + if self.matrix_client is not self.matrix_client_pool.matrix_client: + return + # Ignore events from unauthorized room if room.room_id not in self.config.allowed_rooms: return # Ignore redactions from ourselves - if event.sender == self.matrix_client.user: + if event.sender in self.config.user_ids: return - logger.debug(f"Received event to remove event ID {event.redacts}") + logger.debug( + f"Bot {self.matrix_client.user_id} | Room ID {room.room_id} | " + f"Event ID {event.event_id} | Sender {event.sender} | " + f"Received event to remove event ID {event.redacts}" + ) command = UnackAlertCommand( self.matrix_client, @@ -243,7 +300,12 @@ class Callbacks: try: await command.process() except (SendRetryError, LocalProtocolError) as e: - logger.exception(f"Unable to send message to {room.room_id}", exc_info=e) + logger.exception( + f"Bot {self.matrix_client.user_id} | Room ID {room.room_id} | " + f"Event ID {event.event_id} | Sender {event.sender} | " + f"Cannot send message to room.", + exc_info=e, + ) async def decryption_failure(self, room: MatrixRoom, event: MegolmEvent) -> None: """Callback for when an event fails to decrypt. Inform the user. @@ -258,7 +320,9 @@ class Callbacks: return logger.error( - f"Failed to decrypt event '{event.event_id}' in room '{room.room_id}'!" + f"Bot {self.matrix_client.user_id} | Room ID {room.room_id} | " + f"Event ID {event.event_id} | Sender {event.sender} | " + f"Failed to decrypt event!" f"\n\n" f"Tip: try using a different device ID in your config file and restart." f"\n\n" @@ -271,7 +335,8 @@ class Callbacks: """Callback for when somebody wants to verify our devices.""" if "emoji" not in event.short_authentication_string: logger.error( - f"Unable to use emoji verification with {event.sender} on device {event.from_device}." + f"Bot {self.matrix_client.user_id} | Sender {event.sender} | " + f"Cannot use emoji verification with device {event.from_device}." ) return @@ -280,7 +345,8 @@ class Callbacks: ) if isinstance(event_response, ToDeviceError): logger.error( - f"Unable to start key verification with {event.sender} on device {event.from_device}, got error: {event_response}." + f"Bot {self.matrix_client.user_id} | Sender {event.sender} | " + f"Cannot start key verification with device {event.from_device}, got error: {event_response}." ) return @@ -290,7 +356,8 @@ class Callbacks: event_response = await self.matrix_client.to_device(todevice_msg) if isinstance(event_response, ToDeviceError): logger.error( - f"Unable to share key with {event.sender} on device {event.from_device}, got error: {event_response}." + f"Bot {self.matrix_client.user_id} | Sender {event.sender} | " + f"Cannot share key with device {event.from_device}, got error: {event_response}." ) return @@ -300,7 +367,8 @@ class Callbacks: # here. The SAS flow is already cancelled. # We only need to inform the user. logger.info( - f"Verification has been cancelled by {event.sender} for reason: {event.reason}." + f"Bot {self.matrix_client.user_id} | Sender {event.sender} | " + f"Key verification has been cancelled for reason: {event.reason}." ) async def key_verification_confirm(self, event: KeyVerificationKey): @@ -310,7 +378,8 @@ class Callbacks: alt_text_str = " ".join(alt_text_list) logger.info( - f"Received request to verify emojis from {event.sender}: {emoji_str} ({alt_text_str})" + f"Bot {self.matrix_client.user_id} | Sender {event.sender} | " + f"Received request to verify emojis: {emoji_str} ({alt_text_str})" ) event_response = await self.matrix_client.confirm_short_auth_string( @@ -318,7 +387,8 @@ class Callbacks: ) if isinstance(event_response, ToDeviceError): logger.error( - f"Unable to confirm emoji verification with {event.sender}, got error: {event_response}." + f"Bot {self.matrix_client.user_id} | Sender {event.sender} | " + f"Cannot confirm emoji verification, got error: {event_response}." ) # FIXME: We should allow manual cancel or reject @@ -338,12 +408,13 @@ class Callbacks: # f"Unable to cancel emoji verification with {event.sender}, got error: {event_response}." # ) - async def key_verification_end(self, event: KeyVerificationMac): + async def key_verification_end(self, event: KeyVerificationMac) -> None: try: sas = self.matrix_client.key_verifications[event.transaction_id] except KeyError: logger.error( - f"Unable to find transaction ID {event.transaction_id} sent by {event.sender}" + f"Bot {self.matrix_client.user_id} | Sender {event.sender} | " + f"Cannot find transaction ID {event.transaction_id}" ) return @@ -351,44 +422,54 @@ class Callbacks: todevice_msg = sas.get_mac() except LocalProtocolError as e: # e.g. it might have been cancelled by ourselves - logger.warning(f"Unable to conclude verification with {event.sender}: {e}.") + logger.warning( + f"Bot {self.matrix_client.user_id} | Sender {event.sender} | " + f"Cannot conclude key verification: {e}." + ) return event_response = await self.matrix_client.to_device(todevice_msg) if isinstance(event_response, ToDeviceError): logger.error( - f"Unable to conclude verification with {event.sender}, got error: {event_response}." + f"Bot {self.matrix_client.user_id} | Sender {event.sender} | " + f"Cannot conclude key verification, got error: {event_response}." ) return verified_devices = " ".join(sas.verified_devices) logger.info( - f"Successfully verified devices from {event.sender}: {verified_devices}" + f"Bot {self.matrix_client.user_id} | Sender {event.sender} | " + f"Successfully verified devices: {verified_devices}" ) - async def unknown(self, room: MatrixRoom, event: UnknownEvent) -> None: - """Callback for when an event with a type that is unknown to matrix-nio is received. - Currently this is used for reaction events, which are not yet part of a released - matrix spec (and are thus unknown to nio). - - Args: - room: The room the reaction was sent in. - - event: The event itself. - """ - # Ignore events from unauthorized room - if room.room_id not in self.config.allowed_rooms: + async def unknown_message( + self, room: MatrixRoom, event: RoomMessageUnknown + ) -> None: + event_content = event.source["content"] + if event_content["msgtype"] != "m.key.verification.request": return - if event.type == "m.reaction": - # Get the ID of the event this was a reaction to - relation_dict = event.source.get("content", {}).get("m.relates_to", {}) + if "m.sas.v1" not in event_content["methods"]: + return - reacted_to_id = relation_dict.get("event_id") - if reacted_to_id and relation_dict.get("rel_type") == "m.annotation": - await self._reaction(room, event, reacted_to_id) - return - - logger.debug( - f"Got unknown event with type to {event.type} from {event.sender} in {room.room_id}." + response_event = await self.matrix_client.room_send( + room.room_id, + "m.room.message", + { + "msgtype": "m.key.verification.ready", + "methods": ["m.sas.v1"], + "m.relates_to": {"rel_type": "m.reference", "event_id": event.event_id}, + }, + ) + + if isinstance(response_event, RoomSendError): + raise SendRetryError( + f"{response_event.status_code} - {response_event.message}" + ) + + async def debug(self, room: MatrixRoom, event: Event) -> None: + logger.debug( + f"Bot {self.matrix_client.user_id} | Room ID {room.room_id} | " + f"Event ID {event.event_id} | Sender {event.sender} | " + f"Received some event: {event.source}" ) diff --git a/matrix_alertbot/chat_functions.py b/matrix_alertbot/chat_functions.py index 450f2ee..ab229d0 100644 --- a/matrix_alertbot/chat_functions.py +++ b/matrix_alertbot/chat_functions.py @@ -1,14 +1,11 @@ +from __future__ import annotations + import logging from typing import Dict, Optional, TypedDict, Union -from nio import ( - AsyncClient, - ErrorResponse, - Response, - RoomSendError, - RoomSendResponse, - SendRetryError, -) +from nio.client import AsyncClient +from nio.exceptions import SendRetryError +from nio.responses import ErrorResponse, Response, RoomSendError, RoomSendResponse from typing_extensions import NotRequired logger = logging.getLogger(__name__) @@ -30,7 +27,7 @@ async def send_text_to_room( matrix_client: AsyncClient, room_id: str, plaintext: str, - html: str = None, + html: Optional[str] = None, notice: bool = True, reply_to_event_id: Optional[str] = None, ) -> RoomSendResponse: diff --git a/matrix_alertbot/command.py b/matrix_alertbot/command.py index 0883035..d8d29f0 100644 --- a/matrix_alertbot/command.py +++ b/matrix_alertbot/command.py @@ -1,9 +1,12 @@ +from __future__ import annotations + import logging -from typing import Optional, Tuple +from typing import Optional, Tuple, cast import pytimeparse2 from diskcache import Cache -from nio import AsyncClient, MatrixRoom +from nio.client import AsyncClient +from nio.rooms import MatrixRoom from matrix_alertbot.alertmanager import AlertmanagerClient from matrix_alertbot.chat_functions import send_text_to_room @@ -28,7 +31,7 @@ class BaseCommand: room: MatrixRoom, sender: str, event_id: str, - args: Tuple[str, ...] = None, + args: Tuple[str, ...] = (), ) -> None: """A command made by a user. @@ -77,7 +80,7 @@ class BaseAlertCommand(BaseCommand): sender: str, event_id: str, reacted_to_event_id: str, - args: Tuple[str, ...] = None, + args: Tuple[str, ...] = (), ) -> None: super().__init__( client, cache, alertmanager, config, room, sender, event_id, args @@ -94,7 +97,7 @@ class AckAlertCommand(BaseAlertCommand): duration = " ".join(durations) logger.debug(f"Receiving a command to create a silence for {duration}.") - duration_seconds = pytimeparse2.parse(duration) + duration_seconds = cast(Optional[int], pytimeparse2.parse(duration)) if duration_seconds is None: logger.error(f"Unable to create silence: Invalid duration '{duration}'") await send_text_to_room( @@ -123,17 +126,21 @@ class AckAlertCommand(BaseAlertCommand): f"Reading alert fingerprint for event {self.reacted_to_event_id} from cache" ) try: - alert_fingerprint: str = self.cache[self.reacted_to_event_id] + alert_fingerprint = cast(str, self.cache[self.reacted_to_event_id]) except KeyError: logger.error( f"Cannot find fingerprint for alert event {self.reacted_to_event_id} in cache" ) return + sender_user_name = self.room.user_name(self.sender) + if sender_user_name is None: + sender_user_name = self.sender + try: silence_id = await self.alertmanager_client.create_or_update_silence( alert_fingerprint, - self.room.user_name(self.sender), + sender_user_name, duration_seconds, force=True, ) @@ -173,7 +180,7 @@ class UnackAlertCommand(BaseAlertCommand): f"Reading alert fingerprint for event {self.reacted_to_event_id} from cache." ) try: - alert_fingerprint: str = self.cache[self.reacted_to_event_id] + alert_fingerprint = cast(str, self.cache[self.reacted_to_event_id]) except KeyError: logger.error( f"Cannot find fingerprint for alert event {self.reacted_to_event_id} in cache." @@ -185,7 +192,7 @@ class UnackAlertCommand(BaseAlertCommand): f"Reading silence ID for alert fingerprint {alert_fingerprint} from cache." ) try: - silence_id: str = self.cache[alert_fingerprint] + silence_id = cast(str, self.cache[alert_fingerprint]) except KeyError: logger.error( f"Cannot find silence for alert fingerprint {alert_fingerprint} in cache" diff --git a/matrix_alertbot/config.py b/matrix_alertbot/config.py index 6245c1f..3bfe269 100644 --- a/matrix_alertbot/config.py +++ b/matrix_alertbot/config.py @@ -1,8 +1,10 @@ +from __future__ import annotations + import logging import os import re import sys -from typing import Any, List, Optional +from typing import Any, Dict, List, Optional import pytimeparse2 import yaml @@ -22,6 +24,27 @@ logging.getLogger("peewee").setLevel( DEFAULT_REACTIONS = {"🤫", "😶", "🤐", "🙊", "🔇", "🔕"} +class AccountConfig: + def __init__(self, account: Dict[str, str]) -> None: + self.id: str = account["id"] + if not re.match("@.+:.+", self.id): + raise InvalidConfigError("matrix.user_id must be in the form @name:domain") + + self.password: Optional[str] = account.get("password") + self.token: Optional[str] = account.get("token") + + if self.password is None and self.token is None: + raise RequiredConfigKeyError("Must supply either user token or password") + + self.device_id: Optional[str] = account.get("device_id") + self.token_file: str = account.get("token_file", "token.json") + + self.homeserver_url: str = account["url"] + + def __repr__(self) -> str: + return f"{self.__class__.__name__}({self.id})" + + class Config: """Creates a Config object from a YAML-encoded config file from a given filepath""" @@ -53,17 +76,27 @@ class Config: file_logging_filepath = self._get_cfg( ["logging", "file_logging", "filepath"], default="matrix-alertbot.log" ) + file_logging_log_level = self._get_cfg( + ["logging", "file_logging", "level"], required=False + ) if file_logging_enabled: file_handler = logging.FileHandler(file_logging_filepath) file_handler.setFormatter(formatter) + if file_logging_log_level: + file_handler.setLevel(file_logging_log_level) logger.addHandler(file_handler) console_logging_enabled = self._get_cfg( ["logging", "console_logging", "enabled"], default=True ) + console_logging_log_level = self._get_cfg( + ["logging", "console_logging", "level"], required=False + ) if console_logging_enabled: console_handler = logging.StreamHandler(sys.stdout) console_handler.setFormatter(formatter) + if console_logging_log_level: + console_handler.setLevel(console_logging_log_level) logger.addHandler(console_handler) # Storage setup @@ -91,26 +124,22 @@ class Config: ["alertmanager", "url"], required=True ) - # Matrix bot account setup - self.user_id: str = self._get_cfg(["matrix", "user_id"], required=True) - if not re.match("@.+:.+", self.user_id): - raise InvalidConfigError("matrix.user_id must be in the form @name:domain") - - self.user_password: str = self._get_cfg( - ["matrix", "user_password"], required=False - ) - self.user_token: str = self._get_cfg(["matrix", "user_token"], required=False) - if not self.user_token and not self.user_password: - raise RequiredConfigKeyError("Must supply either user token or password") - - self.device_id: str = self._get_cfg(["matrix", "device_id"], required=False) + # Matrix bot accounts setup + self.accounts: List[AccountConfig] = [] + accounts_dict: list = self._get_cfg(["matrix", "accounts"], required=True) + for i, account_dict in enumerate(accounts_dict): + try: + account = AccountConfig(account_dict) + except KeyError as e: + key_name = e.args[0] + raise RequiredConfigKeyError( + f"Config option matrix.accounts.{i}.{key_name} is required" + ) + self.accounts.append(account) + self.user_ids = {account.id for account in self.accounts} self.device_name: str = self._get_cfg( ["matrix", "device_name"], default="matrix-alertbot" ) - self.user_token_file: str = self._get_cfg( - ["matrix", "user_token_file"], default="token.json" - ) - self.homeserver_url: str = self._get_cfg(["matrix", "url"], required=True) self.allowed_rooms: list = self._get_cfg( ["matrix", "allowed_rooms"], required=True ) diff --git a/matrix_alertbot/errors.py b/matrix_alertbot/errors.py index 9be7555..56ae3a7 100644 --- a/matrix_alertbot/errors.py +++ b/matrix_alertbot/errors.py @@ -63,3 +63,9 @@ class AlertmanagerServerError(AlertmanagerError): """An error encountered with Alertmanager server.""" pass + + +class MatrixClientError(MatrixAlertbotError): + """An error encountered with the Matrix client""" + + pass diff --git a/matrix_alertbot/main.py b/matrix_alertbot/main.py index 28d98fd..a04c337 100644 --- a/matrix_alertbot/main.py +++ b/matrix_alertbot/main.py @@ -1,142 +1,20 @@ #!/usr/bin/env python3 -import asyncio -import json -import logging -import os -import sys -from asyncio import TimeoutError +from __future__ import annotations + +import asyncio +import logging +import sys -from aiohttp import ClientConnectionError, ServerDisconnectedError from diskcache import Cache -from nio import ( - AsyncClient, - AsyncClientConfig, - InviteMemberEvent, - KeyVerificationCancel, - KeyVerificationKey, - KeyVerificationMac, - KeyVerificationStart, - LocalProtocolError, - LoginError, - MegolmEvent, - RedactionEvent, - RoomMessageText, - UnknownEvent, -) from matrix_alertbot.alertmanager import AlertmanagerClient -from matrix_alertbot.callback import Callbacks from matrix_alertbot.config import Config +from matrix_alertbot.matrix import MatrixClientPool from matrix_alertbot.webhook import Webhook logger = logging.getLogger(__name__) -def create_matrix_client(config: Config) -> AsyncClient: - # Configuration options for the AsyncClient - try: - matrix_client_config = AsyncClientConfig( - max_limit_exceeded=5, - max_timeouts=3, - store_sync_tokens=True, - encryption_enabled=True, - ) - except ImportWarning as e: - logger.warning(e) - matrix_client_config = AsyncClientConfig( - max_limit_exceeded=5, - max_timeouts=3, - store_sync_tokens=True, - encryption_enabled=False, - ) - - # Load credentials from a previous session - if os.path.exists(config.user_token_file): - with open(config.user_token_file, "r") as ifd: - credentials = json.load(ifd) - config.user_token = credentials["access_token"] - config.device_id = credentials["device_id"] - - # Initialize the matrix client based on stored credentials - matrix_client = AsyncClient( - config.homeserver_url, - config.user_id, - device_id=config.device_id, - store_path=config.store_dir, - config=matrix_client_config, - ) - - return matrix_client - - -async def start_matrix_client( - matrix_client: AsyncClient, cache: Cache, config: Config -) -> bool: - # Keep trying to reconnect on failure (with some time in-between) - while True: - try: - if config.device_id and config.user_token: - matrix_client.restore_login( - user_id=config.user_id, - device_id=config.device_id, - access_token=config.user_token, - ) - - # Sync encryption keys with the server - if matrix_client.should_upload_keys: - await matrix_client.keys_upload() - else: - # Try to login with the configured username/password - try: - login_response = await matrix_client.login( - password=config.user_password, - device_name=config.device_name, - ) - - # Check if login failed - if type(login_response) == LoginError: - logger.error("Failed to login: %s", login_response.message) - return False - except LocalProtocolError as e: - # There's an edge case here where the user hasn't installed the correct C - # dependencies. In that case, a LocalProtocolError is raised on login. - logger.fatal( - "Failed to login. Have you installed the correct dependencies? " - "https://github.com/poljar/matrix-nio#installation " - "Error: %s", - e, - ) - return False - - # Save user's access token and device ID - # See https://stackoverflow.com/a/45368120 - user_token_fd = os.open( - config.user_token_file, - flags=os.O_CREAT | os.O_WRONLY | os.O_TRUNC, - mode=0o640, - ) - with os.fdopen(user_token_fd, "w") as ofd: - json.dump( - { - "device_id": login_response.device_id, - "access_token": login_response.access_token, - }, - ofd, - ) - - # Login succeeded! - - logger.info(f"Logged in as {config.user_id}") - await matrix_client.sync_forever(timeout=30000, full_state=True) - except (ClientConnectionError, ServerDisconnectedError, TimeoutError): - logger.warning("Unable to connect to homeserver, retrying in 15s...") - - # Sleep so we don't bombard the server with login requests - await asyncio.sleep(15) - finally: - await matrix_client.close() - - def main() -> None: """The first function that is run when starting the bot""" @@ -150,41 +28,22 @@ def main() -> None: # Read the parsed config file and create a Config object config = Config(config_path) - matrix_client = create_matrix_client(config) - # Configure the cache cache = Cache(config.cache_dir) # Configure Alertmanager client alertmanager_client = AlertmanagerClient(config.alertmanager_url, cache) - # Set up event callbacks - callbacks = Callbacks(matrix_client, alertmanager_client, cache, config) - matrix_client.add_event_callback(callbacks.message, (RoomMessageText,)) - matrix_client.add_event_callback( - callbacks.invite_event_filtered_callback, (InviteMemberEvent,) - ) - matrix_client.add_event_callback(callbacks.decryption_failure, (MegolmEvent,)) - matrix_client.add_event_callback(callbacks.unknown, (UnknownEvent,)) - matrix_client.add_event_callback(callbacks.redaction, (RedactionEvent,)) - matrix_client.add_to_device_callback( - callbacks.key_verification_start, (KeyVerificationStart,) - ) - matrix_client.add_to_device_callback( - callbacks.key_verification_cancel, (KeyVerificationCancel,) - ) - matrix_client.add_to_device_callback( - callbacks.key_verification_confirm, (KeyVerificationKey,) - ) - matrix_client.add_to_device_callback( - callbacks.key_verification_end, (KeyVerificationMac,) - ) + # Create matrix clients + matrix_client_pool = MatrixClientPool(alertmanager_client, cache, config) # Configure webhook server - webhook_server = Webhook(matrix_client, alertmanager_client, cache, config) + webhook_server = Webhook(matrix_client_pool, alertmanager_client, cache, config) loop = asyncio.get_event_loop() + loop.create_task(matrix_client_pool.switch_active_client()) loop.create_task(webhook_server.start()) - loop.create_task(start_matrix_client(matrix_client, cache, config)) + for account in config.accounts: + loop.create_task(matrix_client_pool.start(account, config)) try: loop.run_forever() @@ -193,5 +52,5 @@ def main() -> None: finally: loop.run_until_complete(webhook_server.close()) loop.run_until_complete(alertmanager_client.close()) - loop.run_until_complete(matrix_client.close()) + loop.run_until_complete(matrix_client_pool.close()) cache.close() diff --git a/matrix_alertbot/matrix.py b/matrix_alertbot/matrix.py new file mode 100644 index 0000000..a13369d --- /dev/null +++ b/matrix_alertbot/matrix.py @@ -0,0 +1,252 @@ +from __future__ import annotations + +import asyncio +import json +import logging +import os +import random +from typing import Dict, List, Optional, Tuple + +from aiohttp import ClientConnectionError, ServerDisconnectedError +from diskcache import Cache +from nio.client import AsyncClient, AsyncClientConfig +from nio.events import ( + InviteMemberEvent, + KeyVerificationCancel, + KeyVerificationKey, + KeyVerificationMac, + KeyVerificationStart, + MegolmEvent, + ReactionEvent, + RedactionEvent, + RoomMessageText, + RoomMessageUnknown, +) +from nio.exceptions import LocalProtocolError +from nio.responses import LoginError, WhoamiError + +import matrix_alertbot.callback +from matrix_alertbot.alertmanager import AlertmanagerClient +from matrix_alertbot.config import AccountConfig, Config + +logger = logging.getLogger(__name__) + + +class MatrixClientPool: + def __init__( + self, alertmanager_client: AlertmanagerClient, cache: Cache, config: Config + ) -> None: + self._lock = asyncio.Lock() + self._matrix_clients: Dict[AccountConfig, AsyncClient] = {} + self._accounts: List[AccountConfig] = [] + + self._accounts = config.accounts + for account in self._accounts: + matrix_client = self._create_matrix_client( + account, alertmanager_client, cache, config + ) + self._matrix_clients[account] = matrix_client + + self.account = next(iter(self._accounts)) + self.matrix_client = self._matrix_clients[self.account] + + async def switch_active_client( + self, + ) -> Optional[Tuple[AsyncClient, AccountConfig]]: + async with self._lock: + for account in random.sample(self._accounts, len(self._accounts)): + if account is self.account: + continue + + matrix_client = self._matrix_clients[account] + try: + whoami = await matrix_client.whoami() + logged_in = not isinstance(whoami, WhoamiError) + except Exception: + logged_in = False + + if logged_in: + self.account = account + self.matrix_client = matrix_client + + logger.warning( + f"Bot {self.account.id} | Matrix client for homeserver {self.account.homeserver_url} selected as new leader." + ) + + return matrix_client, account + + if self.matrix_client.logged_in: + logger.warning( + f"Bot {self.account.id} | No active Matrix client available, keeping Matrix client for {self.account.homeserver_url} as the leader." + ) + else: + logger.error( + f"Bot {self.account.id} | No active Matrix client connected." + ) + return None + + async def close(self) -> None: + for matrix_client in self._matrix_clients.values(): + await matrix_client.close() + + def _create_matrix_client( + self, + account: AccountConfig, + alertmanager_client: AlertmanagerClient, + cache: Cache, + config: Config, + ) -> AsyncClient: + # Configuration options for the AsyncClient + try: + matrix_client_config = AsyncClientConfig( + max_limit_exceeded=5, + max_timeouts=3, + store_sync_tokens=True, + encryption_enabled=True, + ) + except ImportWarning as e: + logger.warning(e) + matrix_client_config = AsyncClientConfig( + max_limit_exceeded=5, + max_timeouts=3, + store_sync_tokens=True, + encryption_enabled=False, + ) + + # Load credentials from a previous session + if os.path.exists(account.token_file): + with open(account.token_file, "r") as ifd: + credentials = json.load(ifd) + account.token = credentials["access_token"] + account.device_id = credentials["device_id"] + + # Initialize the matrix client based on stored credentials + matrix_client = AsyncClient( + account.homeserver_url, + account.id, + device_id=account.device_id, + store_path=config.store_dir, + config=matrix_client_config, + ) + + # Set up event callbacks + callbacks = matrix_alertbot.callback.Callbacks( + matrix_client, alertmanager_client, cache, config, self + ) + + matrix_client.add_event_callback(callbacks.message, (RoomMessageText,)) + matrix_client.add_event_callback( + callbacks.invite_event_filtered_callback, (InviteMemberEvent,) + ) + # matrix_client.add_event_callback(callbacks.debug, (Event,)) + matrix_client.add_event_callback(callbacks.decryption_failure, (MegolmEvent,)) + matrix_client.add_event_callback(callbacks.reaction, (ReactionEvent,)) + matrix_client.add_event_callback(callbacks.redaction, (RedactionEvent,)) + matrix_client.add_event_callback( + callbacks.unknown_message, (RoomMessageUnknown,) + ) + matrix_client.add_to_device_callback( + callbacks.key_verification_start, (KeyVerificationStart,) + ) + matrix_client.add_to_device_callback( + callbacks.key_verification_cancel, (KeyVerificationCancel,) + ) + matrix_client.add_to_device_callback( + callbacks.key_verification_confirm, (KeyVerificationKey,) + ) + matrix_client.add_to_device_callback( + callbacks.key_verification_end, (KeyVerificationMac,) + ) + + return matrix_client + + async def start( + self, + account: AccountConfig, + config: Config, + ): + matrix_client = self._matrix_clients[account] + + # Keep trying to reconnect on failure (with some time in-between) + # We switch homeserver after some retries + while True: + try: + if account.device_id and account.token: + matrix_client.restore_login( + user_id=account.id, + device_id=account.device_id, + access_token=account.token, + ) + + # Sync encryption keys with the server + if matrix_client.should_upload_keys: + await matrix_client.keys_upload() + else: + # Try to login with the configured username/password + try: + login_response = await matrix_client.login( + password=account.password, + device_name=config.device_name, + ) + + # Check if login failed + if type(login_response) == LoginError: + logger.error( + f"Bot {account.id} | Failed to login: {login_response.message}" + ) + return False + except LocalProtocolError as e: + # There's an edge case here where the user hasn't installed the correct C + # dependencies. In that case, a LocalProtocolError is raised on login. + logger.fatal( + f"Bot {account.id} | Failed to login. Have you installed the correct dependencies? " + "https://github.com/poljar/matrix-nio#installation " + "Error: %s", + e, + ) + return False + + if isinstance(login_response, LoginError): + logger.fatal( + f"Bot {account.id} | Failed to login: {login_response.message}" + ) + return False + + # Save user's access token and device ID + # See https://stackoverflow.com/a/45368120 + account_token_fd = os.open( + account.token_file, + flags=os.O_CREAT | os.O_WRONLY | os.O_TRUNC, + mode=0o640, + ) + with os.fdopen(account_token_fd, "w") as ofd: + json.dump( + { + "device_id": login_response.device_id, + "access_token": login_response.access_token, + }, + ofd, + ) + + # Login succeeded! + + logger.info(f"Bot {account.id} | Logged in.") + + await matrix_client.sync_forever(timeout=30000, full_state=True) + except (ClientConnectionError, ServerDisconnectedError, TimeoutError): + await matrix_client.close() + + logger.warning( + f"Bot {account.id} | Matrix client disconnected, retrying in 15s..." + ) + + if len(self._accounts) > 1 and self.matrix_client is matrix_client: + logger.warning( + f"Bot {account.id} | Selecting another Matrix client as leader..." + ) + await self.switch_active_client() + + # Sleep so we don't bombard the server with login requests + await asyncio.sleep(15) + finally: + await matrix_client.close() diff --git a/matrix_alertbot/webhook.py b/matrix_alertbot/webhook.py index 7443d9a..fac4860 100644 --- a/matrix_alertbot/webhook.py +++ b/matrix_alertbot/webhook.py @@ -1,6 +1,7 @@ from __future__ import annotations import logging +from typing import List import prometheus_client from aiohttp import ClientError, web, web_request @@ -8,7 +9,7 @@ from aiohttp.abc import AbstractAccessLogger from aiohttp_prometheus_exporter.handler import metrics from aiohttp_prometheus_exporter.middleware import prometheus_middleware_factory from diskcache import Cache -from nio import AsyncClient, LocalProtocolError, SendRetryError +from nio.exceptions import LocalProtocolError, SendRetryError from matrix_alertbot.alert import Alert, AlertRenderer from matrix_alertbot.alertmanager import AlertmanagerClient @@ -16,9 +17,11 @@ from matrix_alertbot.chat_functions import send_text_to_room from matrix_alertbot.config import Config from matrix_alertbot.errors import ( AlertmanagerError, + MatrixClientError, SilenceExtendError, SilenceNotFoundError, ) +from matrix_alertbot.matrix import MatrixClientPool logger = logging.getLogger(__name__) @@ -93,13 +96,13 @@ async def create_alerts(request: web_request.Request) -> web.Response: if len(data["alerts"]) == 0: return web.Response(status=400, body="Alerts cannot be empty.") - alerts = [] - for alert in alert_dicts: + alerts: List[Alert] = [] + for alert_dict in alert_dicts: try: - alert = Alert.from_dict(alert) + alert = Alert.from_dict(alert_dict) except KeyError as e: logger.error(f"Cannot parse alert dict: {e}") - return web.Response(status=400, body=f"Invalid alert: {alert}.") + return web.Response(status=400, body=f"Invalid alert: {alert_dict}.") alerts.append(alert) for alert in alerts: @@ -121,6 +124,14 @@ async def create_alerts(request: web_request.Request) -> web.Response: status=500, body=f"An error occured when sending alert with fingerprint '{alert.fingerprint}' to Matrix room.", ) + except MatrixClientError as e: + logger.error( + f"Unable to send alert {alert.fingerprint} to Matrix room {room_id}: {e}" + ) + return web.Response( + status=500, + body=f"An error occured when sending alert with fingerprint '{alert.fingerprint}' to Matrix room.", + ) except Exception as e: logger.error( f"Unable to send alert {alert.fingerprint} to Matrix room {room_id}: {e}" @@ -138,7 +149,7 @@ async def create_alert( ) -> None: alertmanager_client: AlertmanagerClient = request.app["alertmanager_client"] alert_renderer: AlertRenderer = request.app["alert_renderer"] - matrix_client: AsyncClient = request.app["matrix_client"] + matrix_client_pool: MatrixClientPool = request.app["matrix_client_pool"] cache: Cache = request.app["cache"] config: Config = request.app["config"] @@ -162,9 +173,12 @@ async def create_alert( plaintext = alert_renderer.render(alert, html=False) html = alert_renderer.render(alert, html=True) - event = await send_text_to_room( - matrix_client, room_id, plaintext, html, notice=False - ) + if matrix_client_pool.matrix_client is not None: + event = await send_text_to_room( + matrix_client_pool.matrix_client, room_id, plaintext, html, notice=False + ) + else: + raise MatrixClientError("No matrix client available") if alert.firing: cache.set(event.event_id, alert.fingerprint, expire=config.cache_expire_time) @@ -175,13 +189,13 @@ async def create_alert( class Webhook: def __init__( self, - matrix_client: AsyncClient, + matrix_client_pool: MatrixClientPool, alertmanager_client: AlertmanagerClient, cache: Cache, config: Config, ) -> None: self.app = web.Application(logger=logger) - self.app["matrix_client"] = matrix_client + self.app["matrix_client_pool"] = matrix_client_pool self.app["alertmanager_client"] = alertmanager_client self.app["config"] = config self.app["cache"] = cache diff --git a/setup.cfg b/setup.cfg index 9a24810..2dc6b01 100644 --- a/setup.cfg +++ b/setup.cfg @@ -26,7 +26,7 @@ install_requires = aiotools>=1.5.9 diskcache>=5.4.0 jinja2>=3.1.2 - matrix-nio>=0.19.0 + matrix-nio>=0.24.0 Markdown>=3.3.7 pytimeparse2>=1.4.0 PyYAML>=5.4.1 @@ -56,7 +56,7 @@ test = types-PyYAML>=6.0.9 types-setuptools>=62.6.0 e2e = - matrix-nio[e2e]>=0.19.0 + matrix-nio[e2e]>=0.24.0 all = %(test)s %(e2e)s diff --git a/tests/resources/config/config.full.yml b/tests/resources/config/config.full.yml index 9c73f6c..15e8623 100644 --- a/tests/resources/config/config.full.yml +++ b/tests/resources/config/config.full.yml @@ -7,28 +7,52 @@ command_prefix: "!alert" # Options for connecting to the bot's Matrix account matrix: - # The Matrix User ID of the bot account - user_id: "@fakes_user:matrix.example.com" + accounts: + - # The Matrix User ID of the bot account + id: "@fakes_user:matrix.example.com" - # Matrix account password (optional if access token used) - user_password: "password" + # Matrix account password (optional if access token used) + password: "password" - # Matrix account access token (optional if password used) - # If not set, the server will provide an access token after log in, - # which will be stored in the user token file (see below) - #user_token: "" + # Matrix account access token (optional if password used) + # If not set, the server will provide an access token after log in, + # which will be stored in the user token file (see below) + #token: "" - # Path to the file where to store the user access token - user_token_file: "token.json" + # Path to the file where to store the user access token + token_file: "fake_token.json" - # The URL of the homeserver to connect to - url: https://matrix.example.com + # The URL of the homeserver to connect to + url: https://matrix.example.com - # The device ID that is **non pre-existing** device - # If this device ID already exists, messages will be dropped silently in encrypted rooms - # If not set the server will provide a device ID after log in. Note that this ID - # will change each time the bot reconnects. - device_id: ABCDEFGHIJ + # The device ID that is **non pre-existing** device + # If this device ID already exists, messages will be dropped silently in encrypted rooms + # If not set the server will provide a device ID after log in. Note that this ID + # will change each time the bot reconnects. + device_id: ABCDEFGHIJ + + - # The Matrix User ID of the bot account + id: "@other_user:matrix.domain.tld" + + # Matrix account password (optional if access token used) + #password: "password" + + # Matrix account access token (optional if password used) + # If not set, the server will provide an access token after log in, + # which will be stored in the user token file (see below) + token: "token" + + # Path to the file where to store the user access token + token_file: "other_token.json" + + # The URL of the homeserver to connect to + url: https://matrix.domain.tld + + # The device ID that is **non pre-existing** device + # If this device ID already exists, messages will be dropped silently in encrypted rooms + # If not set the server will provide a device ID after log in. Note that this ID + # will change each time the bot reconnects. + device_id: KLMNOPQRST # What to name the logged in device device_name: fake_device_name @@ -75,6 +99,7 @@ logging: enabled: true # The path to the file to log to. May be relative or absolute filepath: fake.log + level: INFO # Configure logging to the console output console_logging: # Whether logging to the console is enabled diff --git a/tests/resources/config/config.minimal.yml b/tests/resources/config/config.minimal.yml index e1634d7..9c3682a 100644 --- a/tests/resources/config/config.minimal.yml +++ b/tests/resources/config/config.minimal.yml @@ -4,14 +4,15 @@ # Options for connecting to the bot's Matrix account matrix: - # The Matrix User ID of the bot account - user_id: "@fakes_user:matrix.example.com" + accounts: + - # The Matrix User ID of the bot account + id: "@fakes_user:matrix.example.com" - # Matrix account password (optional if access token used) - user_password: "password" + # Matrix account password (optional if access token used) + password: "password" - # The URL of the homeserver to connect to - url: https://matrix.example.com + # The URL of the homeserver to connect to + url: https://matrix.example.com # List of rooms where the bot can interact allowed_rooms: diff --git a/tests/test_callback.py b/tests/test_callback.py index 1e46016..d3e288c 100644 --- a/tests/test_callback.py +++ b/tests/test_callback.py @@ -1,3 +1,5 @@ +from __future__ import annotations + import unittest from typing import Dict from unittest.mock import MagicMock, Mock, patch @@ -6,13 +8,10 @@ import nio import nio.crypto from diskcache import Cache +import matrix_alertbot.alertmanager import matrix_alertbot.callback import matrix_alertbot.command -from matrix_alertbot.alertmanager import AlertmanagerClient -from matrix_alertbot.callback import Callbacks -from matrix_alertbot.command import BaseCommand - -from tests.utils import make_awaitable +import matrix_alertbot.matrix def key_verification_get_mac_raise_protocol_error(): @@ -23,10 +22,13 @@ class CallbacksTestCase(unittest.IsolatedAsyncioTestCase): def setUp(self) -> None: # Create a Callbacks object and give it some Mock'd objects to use self.fake_matrix_client = Mock(spec=nio.AsyncClient) - self.fake_matrix_client.user = "@fake_user:example.com" + self.fake_matrix_client.user_id = "@fake_user:example.com" + # self.fake_matrix_client.user = "@fake_user" self.fake_cache = MagicMock(spec=Cache) - self.fake_alertmanager_client = Mock(spec=AlertmanagerClient) + self.fake_alertmanager_client = Mock( + spec=matrix_alertbot.alertmanager.AlertmanagerClient + ) # Create a fake room to play with self.fake_room = Mock(spec=nio.MatrixRoom) @@ -38,12 +40,19 @@ class CallbacksTestCase(unittest.IsolatedAsyncioTestCase): self.fake_config.allowed_rooms = [self.fake_room.room_id] self.fake_config.allowed_reactions = ["🤫"] self.fake_config.command_prefix = "!alert " + self.fake_config.user_ids = [self.fake_matrix_client.user_id] - self.callbacks = Callbacks( + self.fake_matrix_client_pool = Mock( + spec=matrix_alertbot.matrix.MatrixClientPool + ) + self.fake_matrix_client_pool.matrix_client = self.fake_matrix_client + + self.callbacks = matrix_alertbot.callback.Callbacks( self.fake_matrix_client, self.fake_alertmanager_client, self.fake_cache, self.fake_config, + self.fake_matrix_client_pool, ) async def test_invite(self) -> None: @@ -52,9 +61,6 @@ class CallbacksTestCase(unittest.IsolatedAsyncioTestCase): fake_invite_event = Mock(spec=nio.InviteMemberEvent) fake_invite_event.sender = "@some_other_fake_user:example.com" - # Pretend that attempting to join a room is always successful - self.fake_matrix_client.join.return_value = make_awaitable() - # Pretend that we received an invite event await self.callbacks.invite(self.fake_room, fake_invite_event) @@ -68,6 +74,7 @@ class CallbacksTestCase(unittest.IsolatedAsyncioTestCase): fake_message_event = Mock(spec=nio.RoomMessageText) fake_message_event.sender = "@some_other_fake_user:example.com" fake_message_event.body = "Hello world!" + fake_message_event.event_id = "some event id" # Pretend that we received a text message event await self.callbacks.message(self.fake_room, fake_message_event) @@ -140,7 +147,7 @@ class CallbacksTestCase(unittest.IsolatedAsyncioTestCase): # Tests that the bot process messages in the room that contain a command fake_message_event = Mock(spec=nio.RoomMessageText) - fake_message_event.sender = self.fake_matrix_client.user + fake_message_event.sender = self.fake_matrix_client.user_id # Pretend that we received a text message event await self.callbacks.message(self.fake_room, fake_message_event) @@ -270,30 +277,20 @@ class CallbacksTestCase(unittest.IsolatedAsyncioTestCase): # Tests that the bot process messages in the room that contain a command fake_alert_event = Mock(spec=nio.RoomMessageText) fake_alert_event.event_id = "some alert event id" - fake_alert_event.sender = self.fake_config.user_id + fake_alert_event.sender = self.fake_matrix_client.user_id - fake_reaction_event = Mock(spec=nio.UnknownEvent) - fake_reaction_event.type = "m.reaction" + fake_reaction_event = Mock(spec=nio.ReactionEvent) fake_reaction_event.event_id = "some event id" fake_reaction_event.sender = "@some_other_fake_user:example.com" - fake_reaction_event.source = { - "content": { - "m.relates_to": { - "event_id": fake_alert_event.event_id, - "key": "🤫", - "rel_type": "m.annotation", - } - } - } + fake_reaction_event.reacts_to = fake_alert_event.event_id + fake_reaction_event.key = "🤫" fake_event_response = Mock(spec=nio.RoomGetEventResponse) fake_event_response.event = fake_alert_event - self.fake_matrix_client.room_get_event.return_value = make_awaitable( - fake_event_response - ) + self.fake_matrix_client.room_get_event.return_value = fake_event_response # Pretend that we received a text message event - await self.callbacks.unknown(self.fake_room, fake_reaction_event) + await self.callbacks.reaction(self.fake_room, fake_reaction_event) # Check that we attempted to execute the command fake_command.assert_called_once_with( @@ -317,27 +314,18 @@ class CallbacksTestCase(unittest.IsolatedAsyncioTestCase): # Tests that the bot process messages in the room that contain a command fake_alert_event_id = "some alert event id" - fake_reaction_event = Mock(spec=nio.UnknownEvent) + fake_reaction_event = Mock(spec=nio.ReactionEvent) fake_reaction_event.type = "m.reaction" fake_reaction_event.event_id = "some event id" fake_reaction_event.sender = "@some_other_fake_user:example.com" - fake_reaction_event.source = { - "content": { - "m.relates_to": { - "event_id": fake_alert_event_id, - "key": "🤫", - "rel_type": "m.annotation", - } - } - } + fake_reaction_event.reacts_to = fake_alert_event_id + fake_reaction_event.key = "🤫" fake_event_response = Mock(spec=nio.RoomGetEventError) - self.fake_matrix_client.room_get_event.return_value = make_awaitable( - fake_event_response - ) + self.fake_matrix_client.room_get_event.return_value = fake_event_response # Pretend that we received a text message event - await self.callbacks.unknown(self.fake_room, fake_reaction_event) + await self.callbacks.reaction(self.fake_room, fake_reaction_event) # Check that we attempted to execute the command fake_command.assert_not_called() @@ -356,28 +344,19 @@ class CallbacksTestCase(unittest.IsolatedAsyncioTestCase): fake_alert_event.event_id = "some alert event id" fake_alert_event.sender = "@some_other_fake_user.example.com" - fake_reaction_event = Mock(spec=nio.UnknownEvent) + fake_reaction_event = Mock(spec=nio.ReactionEvent) fake_reaction_event.type = "m.reaction" fake_reaction_event.event_id = "some event id" fake_reaction_event.sender = "@some_other_fake_user:example.com" - fake_reaction_event.source = { - "content": { - "m.relates_to": { - "event_id": fake_alert_event.event_id, - "key": "🤫", - "rel_type": "m.annotation", - } - } - } + fake_reaction_event.reacts_to = fake_alert_event.event_id + fake_reaction_event.key = "🤫" fake_event_response = Mock(spec=nio.RoomGetEventResponse) fake_event_response.event = fake_alert_event - self.fake_matrix_client.room_get_event.return_value = make_awaitable( - fake_event_response - ) + self.fake_matrix_client.room_get_event.return_value = fake_event_response # Pretend that we received a text message event - await self.callbacks.unknown(self.fake_room, fake_reaction_event) + await self.callbacks.reaction(self.fake_room, fake_reaction_event) # Check that we attempted to execute the command fake_command.assert_not_called() @@ -392,22 +371,15 @@ class CallbacksTestCase(unittest.IsolatedAsyncioTestCase): # Tests that the bot process messages in the room that contain a command fake_alert_event_id = "some alert event id" - fake_reaction_event = Mock(spec=nio.UnknownEvent) + fake_reaction_event = Mock(spec=nio.ReactionEvent) fake_reaction_event.type = "m.reaction" fake_reaction_event.event_id = "some event id" fake_reaction_event.sender = "@some_other_fake_user:example.com" - fake_reaction_event.source = { - "content": { - "m.relates_to": { - "event_id": fake_alert_event_id, - "key": "unknown", - "rel_type": "m.annotation", - } - } - } + fake_reaction_event.reacts_to = fake_alert_event_id + fake_reaction_event.key = "unknown" # Pretend that we received a text message event - await self.callbacks.unknown(self.fake_room, fake_reaction_event) + await self.callbacks.reaction(self.fake_room, fake_reaction_event) # Check that we attempted to execute the command fake_command.assert_not_called() @@ -419,25 +391,15 @@ class CallbacksTestCase(unittest.IsolatedAsyncioTestCase): # Tests that the bot process messages in the room that contain a command fake_alert_event_id = "some alert event id" - fake_reaction_event = Mock(spec=nio.UnknownEvent) + fake_reaction_event = Mock(spec=nio.ReactionEvent) fake_reaction_event.type = "m.reaction" fake_reaction_event.event_id = "some event id" - fake_reaction_event.sender = self.fake_matrix_client.user - fake_reaction_event.source = { - "content": { - "m.relates_to": { - "event_id": fake_alert_event_id, - "key": "unknown", - "rel_type": "m.annotation", - } - } - } + fake_reaction_event.sender = self.fake_matrix_client.user_id + fake_reaction_event.reacts_to = fake_alert_event_id + fake_reaction_event.key = "unknown" # Pretend that we received a text message event - await self.callbacks.unknown(self.fake_room, fake_reaction_event) - await self.callbacks._reaction( - self.fake_room, fake_reaction_event, fake_alert_event_id - ) + await self.callbacks.reaction(self.fake_room, fake_reaction_event) # Check that we attempted to execute the command fake_command.assert_not_called() @@ -453,25 +415,15 @@ class CallbacksTestCase(unittest.IsolatedAsyncioTestCase): fake_alert_event_id = "some alert event id" - fake_reaction_event = Mock(spec=nio.UnknownEvent) + fake_reaction_event = Mock(spec=nio.ReactionEvent) fake_reaction_event.type = "m.reaction" fake_reaction_event.event_id = "some event id" fake_reaction_event.sender = "@some_other_fake_user:example.com" - fake_reaction_event.source = { - "content": { - "m.relates_to": { - "event_id": fake_alert_event_id, - "key": "unknown", - "rel_type": "m.annotation", - } - } - } + fake_reaction_event.reacts_to = fake_alert_event_id + fake_reaction_event.key = "unknown" # Pretend that we received a text message event - await self.callbacks.unknown(self.fake_room, fake_reaction_event) - await self.callbacks._reaction( - self.fake_room, fake_reaction_event, fake_alert_event_id - ) + await self.callbacks.reaction(self.fake_room, fake_reaction_event) # Check that we attempted to execute the command fake_command.assert_not_called() @@ -512,7 +464,7 @@ class CallbacksTestCase(unittest.IsolatedAsyncioTestCase): """Tests the callback for RoomMessageText with the command prefix""" # Tests that the bot process messages in the room that contain a command fake_redaction_event = Mock(spec=nio.RedactionEvent) - fake_redaction_event.sender = self.fake_matrix_client.user + fake_redaction_event.sender = self.fake_matrix_client.user_id fake_cache_dict: Dict = {} self.fake_cache.__getitem__.side_effect = fake_cache_dict.__getitem__ @@ -556,9 +508,6 @@ class CallbacksTestCase(unittest.IsolatedAsyncioTestCase): fake_key_verification_event.short_authentication_string = ["emoji"] fake_key_verification_event.transaction_id = fake_transaction_id - self.fake_matrix_client.accept_key_verification.return_value = make_awaitable() - self.fake_matrix_client.to_device.return_value = make_awaitable() - fake_sas = Mock() fake_transactions_dict = {fake_transaction_id: fake_sas} self.fake_matrix_client.key_verifications = fake_transactions_dict @@ -583,9 +532,6 @@ class CallbacksTestCase(unittest.IsolatedAsyncioTestCase): fake_key_verification_event.short_authentication_string = [] fake_key_verification_event.transaction_id = fake_transaction_id - self.fake_matrix_client.accept_key_verification.return_value = make_awaitable() - self.fake_matrix_client.to_device.return_value = make_awaitable() - fake_sas = Mock() fake_transactions_dict = {fake_transaction_id: fake_sas} self.fake_matrix_client.key_verifications = fake_transactions_dict @@ -610,10 +556,9 @@ class CallbacksTestCase(unittest.IsolatedAsyncioTestCase): fake_key_verification_event.short_authentication_string = ["emoji"] fake_key_verification_event.transaction_id = fake_transaction_id - self.fake_matrix_client.accept_key_verification.return_value = make_awaitable( - Mock(spec=nio.ToDeviceError) + self.fake_matrix_client.accept_key_verification.return_value = Mock( + spec=nio.ToDeviceError ) - self.fake_matrix_client.to_device.return_value = make_awaitable() fake_sas = Mock() fake_transactions_dict = {fake_transaction_id: fake_sas} @@ -641,10 +586,7 @@ class CallbacksTestCase(unittest.IsolatedAsyncioTestCase): fake_key_verification_event.short_authentication_string = ["emoji"] fake_key_verification_event.transaction_id = fake_transaction_id - self.fake_matrix_client.accept_key_verification.return_value = make_awaitable() - self.fake_matrix_client.to_device.return_value = make_awaitable( - Mock(spec=nio.ToDeviceError) - ) + self.fake_matrix_client.to_device.return_value = Mock(spec=nio.ToDeviceError) fake_sas = Mock() fake_transactions_dict = {fake_transaction_id: fake_sas} @@ -680,10 +622,6 @@ class CallbacksTestCase(unittest.IsolatedAsyncioTestCase): fake_key_verification_event.sender = "@some_other_fake_user:example.com" fake_key_verification_event.transaction_id = fake_transaction_id - self.fake_matrix_client.confirm_short_auth_string.return_value = ( - make_awaitable() - ) - fake_sas = Mock() fake_sas.get_emoji.return_value = [ ("emoji1", "alt text1"), @@ -709,8 +647,8 @@ class CallbacksTestCase(unittest.IsolatedAsyncioTestCase): fake_key_verification_event.sender = "@some_other_fake_user:example.com" fake_key_verification_event.transaction_id = fake_transaction_id - self.fake_matrix_client.confirm_short_auth_string.return_value = make_awaitable( - Mock(spec=nio.ToDeviceError) + self.fake_matrix_client.confirm_short_auth_string.return_value = Mock( + spec=nio.ToDeviceError ) fake_sas = Mock() @@ -738,8 +676,6 @@ class CallbacksTestCase(unittest.IsolatedAsyncioTestCase): fake_key_verification_event.sender = "@some_other_fake_user:example.com" fake_key_verification_event.transaction_id = fake_transaction_id - self.fake_matrix_client.to_device.return_value = make_awaitable() - fake_sas = Mock() fake_sas.verified_devices = ["HGFEDCBA"] fake_transactions_dict = {fake_transaction_id: fake_sas} @@ -761,8 +697,6 @@ class CallbacksTestCase(unittest.IsolatedAsyncioTestCase): fake_key_verification_event.sender = "@some_other_fake_user:example.com" fake_key_verification_event.transaction_id = fake_transaction_id - self.fake_matrix_client.to_device.return_value = make_awaitable() - fake_sas = Mock() fake_transactions_dict = {} self.fake_matrix_client.key_verifications = fake_transactions_dict @@ -783,8 +717,6 @@ class CallbacksTestCase(unittest.IsolatedAsyncioTestCase): fake_key_verification_event.sender = "@some_other_fake_user:example.com" fake_key_verification_event.transaction_id = fake_transaction_id - self.fake_matrix_client.to_device.return_value = make_awaitable() - fake_sas = Mock() fake_sas.get_mac.side_effect = key_verification_get_mac_raise_protocol_error fake_transactions_dict = {fake_transaction_id: fake_sas} @@ -806,9 +738,7 @@ class CallbacksTestCase(unittest.IsolatedAsyncioTestCase): fake_key_verification_event.sender = "@some_other_fake_user:example.com" fake_key_verification_event.transaction_id = fake_transaction_id - self.fake_matrix_client.to_device.return_value = make_awaitable( - Mock(spec=nio.ToDeviceError) - ) + self.fake_matrix_client.to_device.return_value = Mock(spec=nio.ToDeviceError) fake_sas = Mock() fake_transactions_dict = {fake_transaction_id: fake_sas} @@ -821,25 +751,6 @@ class CallbacksTestCase(unittest.IsolatedAsyncioTestCase): fake_sas.get_mac.assert_called_once_with() self.fake_matrix_client.to_device.assert_called_once_with(fake_sas.get_mac()) - @patch.object(matrix_alertbot.callback.CommandFactory, "create", autospec=True) - async def test_unknown(self, fake_command_create: Mock) -> None: - """Tests the callback for RoomMessageText with the command prefix""" - # Tests that the bot process messages in the room that contain a command - fake_command = Mock(spec=BaseCommand) - fake_command_create.return_value = fake_command - - fake_reaction_event = Mock(spec=nio.UnknownEvent) - fake_reaction_event.type = "m.reaction" - fake_reaction_event.event_id = "some event id" - fake_reaction_event.sender = "@some_other_fake_user:example.com" - fake_reaction_event.source = {} - - # Pretend that we received a text message event - await self.callbacks.unknown(self.fake_room, fake_reaction_event) - - # Check that we attempted to execute the command - fake_command_create.assert_not_called() - if __name__ == "__main__": unittest.main() diff --git a/tests/test_chat_functions.py b/tests/test_chat_functions.py index e9d82a0..93732e9 100644 --- a/tests/test_chat_functions.py +++ b/tests/test_chat_functions.py @@ -10,8 +10,6 @@ from matrix_alertbot.chat_functions import ( strip_fallback, ) -from tests.utils import make_awaitable - async def send_room_raise_send_retry_error( room_id: str, @@ -39,11 +37,12 @@ class ChatFunctionsTestCase(unittest.IsolatedAsyncioTestCase): async def test_react_to_event(self) -> None: fake_response = Mock(spec=nio.RoomSendResponse) fake_matrix_client = Mock(spec=nio.AsyncClient) - fake_matrix_client.room_send = Mock(return_value=make_awaitable(fake_response)) fake_room_id = "!abcdefgh:example.com" fake_event_id = "some event id" fake_reaction_text = "some reaction" + fake_matrix_client.room_send.return_value = fake_response + response = await react_to_event( fake_matrix_client, fake_room_id, fake_event_id, fake_reaction_text ) @@ -67,7 +66,7 @@ class ChatFunctionsTestCase(unittest.IsolatedAsyncioTestCase): fake_response.message = "some error" fake_response.status_code = "some status code" fake_matrix_client = Mock(spec=nio.AsyncClient) - fake_matrix_client.room_send.return_value = make_awaitable(fake_response) + fake_matrix_client.room_send.return_value = fake_response fake_room_id = "!abcdefgh:example.com" fake_event_id = "some event id" fake_reaction_text = "some reaction" @@ -93,11 +92,12 @@ class ChatFunctionsTestCase(unittest.IsolatedAsyncioTestCase): async def test_send_text_to_room_as_notice(self) -> None: fake_response = Mock(spec=nio.RoomSendResponse) fake_matrix_client = Mock(spec=nio.AsyncClient) - fake_matrix_client.room_send = Mock(return_value=make_awaitable(fake_response)) fake_room_id = "!abcdefgh:example.com" fake_plaintext_body = "some plaintext message" fake_html_body = "some html message" + fake_matrix_client.room_send.return_value = fake_response + response = await send_text_to_room( fake_matrix_client, fake_room_id, fake_plaintext_body, fake_html_body ) @@ -118,11 +118,12 @@ class ChatFunctionsTestCase(unittest.IsolatedAsyncioTestCase): async def test_send_text_to_room_as_message(self) -> None: fake_response = Mock(spec=nio.RoomSendResponse) fake_matrix_client = Mock(spec=nio.AsyncClient) - fake_matrix_client.room_send.return_value = make_awaitable(fake_response) fake_room_id = "!abcdefgh:example.com" fake_plaintext_body = "some plaintext message" fake_html_body = "some html message" + fake_matrix_client.room_send.return_value = fake_response + response = await send_text_to_room( fake_matrix_client, fake_room_id, @@ -147,12 +148,13 @@ class ChatFunctionsTestCase(unittest.IsolatedAsyncioTestCase): async def test_send_text_to_room_in_reply_to_event(self) -> None: fake_response = Mock(spec=nio.RoomSendResponse) fake_matrix_client = Mock(spec=nio.AsyncClient) - fake_matrix_client.room_send.return_value = make_awaitable(fake_response) fake_room_id = "!abcdefgh:example.com" fake_plaintext_body = "some plaintext message" fake_html_body = "some html message" fake_event_id = "some event id" + fake_matrix_client.room_send.return_value = fake_response + response = await send_text_to_room( fake_matrix_client, fake_room_id, @@ -208,11 +210,12 @@ class ChatFunctionsTestCase(unittest.IsolatedAsyncioTestCase): fake_response.status_code = "some status_code" fake_response.message = "some error" fake_matrix_client = Mock(spec=nio.AsyncClient) - fake_matrix_client.room_send.return_value = make_awaitable(fake_response) fake_room_id = "!abcdefgh:example.com" fake_plaintext_body = "some plaintext message" fake_html_body = "some html message" + fake_matrix_client.room_send.return_value = fake_response + with self.assertRaises(nio.SendRetryError): await send_text_to_room( fake_matrix_client, diff --git a/tests/test_command.py b/tests/test_command.py index a298ebe..b7db170 100644 --- a/tests/test_command.py +++ b/tests/test_command.py @@ -21,8 +21,6 @@ from matrix_alertbot.errors import ( SilenceNotFoundError, ) -from tests.utils import make_awaitable - def cache_get_item(key: str) -> str: return { @@ -84,8 +82,6 @@ class CommandTestCase(unittest.IsolatedAsyncioTestCase): # Create a Command object and give it some Mock'd objects to use self.fake_matrix_client = Mock(spec=nio.AsyncClient) self.fake_matrix_client.user = "@fake_user:example.com" - # Pretend that attempting to send a message is always successful - self.fake_matrix_client.room_send.return_value = make_awaitable() self.fake_cache = MagicMock(spec=Cache) self.fake_cache.__getitem__.side_effect = cache_get_item diff --git a/tests/test_config.py b/tests/test_config.py index 425239e..8ef0f68 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -51,12 +51,15 @@ class ConfigTestCase(unittest.TestCase): fake_path_exists.assert_called_once_with("data/store") fake_mkdir.assert_called_once_with("data/store") - self.assertEqual("@fakes_user:matrix.example.com", config.user_id) - self.assertEqual("password", config.user_password) - self.assertIsNone(config.user_token) - self.assertIsNone(config.device_id) + self.assertEqual({"@fakes_user:matrix.example.com"}, config.user_ids) + self.assertEqual(1, len(config.accounts)) + self.assertEqual("password", config.accounts[0].password) + self.assertIsNone(config.accounts[0].token) + self.assertIsNone(config.accounts[0].device_id) self.assertEqual("matrix-alertbot", config.device_name) - self.assertEqual("https://matrix.example.com", config.homeserver_url) + self.assertEqual( + "https://matrix.example.com", config.accounts[0].homeserver_url + ) self.assertEqual(["!abcdefgh:matrix.example.com"], config.allowed_rooms) self.assertEqual(DEFAULT_REACTIONS, config.allowed_reactions) @@ -92,13 +95,24 @@ class ConfigTestCase(unittest.TestCase): fake_path_exists.assert_called_once_with("data/store") fake_mkdir.assert_called_once_with("data/store") - self.assertEqual("@fakes_user:matrix.example.com", config.user_id) - self.assertEqual("password", config.user_password) - self.assertIsNone(config.user_token) - self.assertEqual("token.json", config.user_token_file) - self.assertEqual("ABCDEFGHIJ", config.device_id) + self.assertEqual( + {"@fakes_user:matrix.example.com", "@other_user:matrix.domain.tld"}, + config.user_ids, + ) + self.assertEqual(2, len(config.accounts)) + self.assertEqual("password", config.accounts[0].password) + self.assertIsNone(config.accounts[0].token) + self.assertEqual("fake_token.json", config.accounts[0].token_file) + self.assertEqual("ABCDEFGHIJ", config.accounts[0].device_id) + self.assertEqual( + "https://matrix.example.com", config.accounts[0].homeserver_url + ) + self.assertIsNone(config.accounts[1].password) + self.assertEqual("token", config.accounts[1].token) + self.assertEqual("other_token.json", config.accounts[1].token_file) + self.assertEqual("KLMNOPQRST", config.accounts[1].device_id) + self.assertEqual("https://matrix.domain.tld", config.accounts[1].homeserver_url) self.assertEqual("fake_device_name", config.device_name) - self.assertEqual("https://matrix.example.com", config.homeserver_url) self.assertEqual(["!abcdefgh:matrix.example.com"], config.allowed_rooms) self.assertEqual({"🤫", "😶", "🤐"}, config.allowed_reactions) @@ -150,7 +164,7 @@ class ConfigTestCase(unittest.TestCase): config_path = os.path.join(CONFIG_RESOURCES_DIR, "config.minimal.yml") config = DummyConfig(config_path) - del config.config_dict["matrix"]["user_id"] + del config.config_dict["matrix"]["accounts"] with self.assertRaises(RequiredConfigKeyError): config._parse_config_values() @@ -166,7 +180,7 @@ class ConfigTestCase(unittest.TestCase): config_path = os.path.join(CONFIG_RESOURCES_DIR, "config.minimal.yml") config = DummyConfig(config_path) - del config.config_dict["matrix"]["user_password"] + del config.config_dict["matrix"]["accounts"][0]["password"] with self.assertRaises(RequiredConfigKeyError): config._parse_config_values() @@ -182,7 +196,7 @@ class ConfigTestCase(unittest.TestCase): config_path = os.path.join(CONFIG_RESOURCES_DIR, "config.minimal.yml") config = DummyConfig(config_path) - del config.config_dict["matrix"]["url"] + del config.config_dict["matrix"]["accounts"][0]["url"] with self.assertRaises(RequiredConfigKeyError): config._parse_config_values() @@ -279,27 +293,27 @@ class ConfigTestCase(unittest.TestCase): config_path = os.path.join(CONFIG_RESOURCES_DIR, "config.minimal.yml") config = DummyConfig(config_path) - config.config_dict["matrix"]["user_id"] = "" + config.config_dict["matrix"]["accounts"][0]["id"] = "" with self.assertRaises(InvalidConfigError): config._parse_config_values() - config.config_dict["matrix"]["user_id"] = "@fake_user" + config.config_dict["matrix"]["accounts"][0]["id"] = "@fake_user" with self.assertRaises(InvalidConfigError): config._parse_config_values() - config.config_dict["matrix"]["user_id"] = "@fake_user:" + config.config_dict["matrix"]["accounts"][0]["id"] = "@fake_user:" with self.assertRaises(InvalidConfigError): config._parse_config_values() - config.config_dict["matrix"]["user_id"] = ":matrix.example.com" + config.config_dict["matrix"]["accounts"][0]["id"] = ":matrix.example.com" with self.assertRaises(InvalidConfigError): config._parse_config_values() - config.config_dict["matrix"]["user_id"] = "@:matrix.example.com" + config.config_dict["matrix"]["accounts"][0]["id"] = "@:matrix.example.com" with self.assertRaises(InvalidConfigError): config._parse_config_values() - config.config_dict["matrix"]["user_id"] = "@:" + config.config_dict["matrix"]["accounts"][0]["id"] = "@:" with self.assertRaises(InvalidConfigError): config._parse_config_values() diff --git a/tests/test_webhook.py b/tests/test_webhook.py index 46ee530..8713377 100644 --- a/tests/test_webhook.py +++ b/tests/test_webhook.py @@ -6,7 +6,8 @@ import aiohttp.test_utils import nio from aiohttp import web from diskcache import Cache -from nio import LocalProtocolError, RoomSendResponse +from nio.exceptions import LocalProtocolError +from nio.responses import RoomSendResponse import matrix_alertbot.webhook from matrix_alertbot.alertmanager import AlertmanagerClient @@ -16,6 +17,7 @@ from matrix_alertbot.errors import ( SilenceExtendError, SilenceNotFoundError, ) +from matrix_alertbot.matrix import MatrixClientPool from matrix_alertbot.webhook import Webhook @@ -40,6 +42,8 @@ def update_silence_raise_alertmanager_error(fingerprint: str) -> str: class WebhookApplicationTestCase(aiohttp.test_utils.AioHTTPTestCase): async def get_application(self) -> web.Application: self.fake_matrix_client = Mock(spec=nio.AsyncClient) + self.fake_matrix_client_pool = Mock(spec=MatrixClientPool) + self.fake_matrix_client_pool.matrix_client = self.fake_matrix_client self.fake_alertmanager_client = Mock(spec=AlertmanagerClient) self.fake_cache = Mock(spec=Cache) @@ -81,7 +85,7 @@ class WebhookApplicationTestCase(aiohttp.test_utils.AioHTTPTestCase): } webhook = Webhook( - self.fake_matrix_client, + self.fake_matrix_client_pool, self.fake_alertmanager_client, self.fake_cache, self.fake_config, diff --git a/tests/utils.py b/tests/utils.py deleted file mode 100644 index 202d0e8..0000000 --- a/tests/utils.py +++ /dev/null @@ -1,22 +0,0 @@ -# Utility functions to make testing easier -import asyncio -from typing import Any, Awaitable - - -def run_coroutine(result: Awaitable[Any]) -> Any: - """Wrapper for asyncio functions to allow them to be run from synchronous functions""" - loop = asyncio.get_event_loop() - result = loop.run_until_complete(result) - loop.close() - return result - - -def make_awaitable(result: Any = None) -> Awaitable[Any]: - """ - Makes an awaitable, suitable for mocking an `async` function. - This uses Futures as they can be awaited multiple times so can be returned - to multiple callers. - """ - future = asyncio.Future() # type: ignore - future.set_result(result) - return future