diff --git a/matrix_alertbot/alertmanager.py b/matrix_alertbot/alertmanager.py index 9cab57a..342d853 100644 --- a/matrix_alertbot/alertmanager.py +++ b/matrix_alertbot/alertmanager.py @@ -74,7 +74,12 @@ class AlertmanagerClient: fingerprint, silence_matchers, user, duration_seconds ) - async def update_silence(self, fingerprint: str) -> str: + async def update_silence( + self, + fingerprint: str, + user: Optional[str] = None, + duration_seconds: Optional[int] = None, + ) -> str: try: silence_id: Optional[str] expire_time: Optional[int] @@ -87,16 +92,29 @@ class AlertmanagerClient: f"Cannot find silence for alert with fingerprint {fingerprint} in cache." ) - if expire_time is not None: - raise SilenceExtendError( - f"Cannot extend silence ID {silence_id} with static duration." - ) + if duration_seconds is None: + if expire_time is not None: + raise SilenceExtendError( + f"Cannot extend silence ID {silence_id} with static duration." + ) silence = await self.get_silence(silence_id) - user = silence["createdBy"] + if user is None: + user = silence["createdBy"] silence_matchers = silence["matchers"] - return await self._create_or_update_silence(fingerprint, silence_matchers, user) + return await self._create_or_update_silence( + fingerprint, silence_matchers, user, duration_seconds + ) + + async def create_or_update_silence( + self, fingerprint: str, user: str, duration_seconds: Optional[int] = None + ) -> str: + try: + silence_id = await self.update_silence(fingerprint, user, duration_seconds) + except SilenceNotFoundError: + silence_id = await self.create_silence(fingerprint, user, duration_seconds) + return silence_id async def _create_or_update_silence( self, diff --git a/matrix_alertbot/command.py b/matrix_alertbot/command.py index a4c23d7..c34ebc8 100644 --- a/matrix_alertbot/command.py +++ b/matrix_alertbot/command.py @@ -131,8 +131,10 @@ class AckAlertCommand(BaseAlertCommand): return try: - silence_id = await self.alertmanager_client.create_silence( - alert_fingerprint, self.room.user_name(self.sender), duration_seconds + silence_id = await self.alertmanager_client.create_or_update_silence( + alert_fingerprint, + self.room.user_name(self.sender), + duration_seconds, ) except AlertNotFoundError as e: logger.warning(f"Unable to create silence: {e}") diff --git a/tests/test_alertmanager.py b/tests/test_alertmanager.py index d125266..9a9b6b6 100644 --- a/tests/test_alertmanager.py +++ b/tests/test_alertmanager.py @@ -4,7 +4,7 @@ import json import unittest from datetime import datetime, timedelta from typing import Any, Dict, Optional, Tuple -from unittest.mock import Mock +from unittest.mock import Mock, patch import aiohttp import aiohttp.test_utils @@ -13,6 +13,7 @@ from aiohttp import web, web_request from diskcache import Cache from freezegun import freeze_time +import matrix_alertbot.alertmanager from matrix_alertbot.alertmanager import AlertmanagerClient from matrix_alertbot.errors import ( AlertmanagerServerError, @@ -23,6 +24,12 @@ from matrix_alertbot.errors import ( ) +async def update_silence_raise_silence_not_found( + fingerprint: str, user: str, duration_seconds: int +) -> str: + raise SilenceNotFoundError + + class FakeCache: def __init__(self, cache_dict: Optional[Dict] = None) -> None: if cache_dict is None: @@ -448,6 +455,81 @@ class AlertmanagerClientTestCase(unittest.IsolatedAsyncioTestCase): await alertmanager_client.get_silence("silence2") self.assertEqual({"fingerprint1": ("silence1", 86400)}, fake_cache.cache) + @freeze_time(datetime.utcfromtimestamp(0)) + async def test_update_silence_override_user_and_duration(self) -> None: + fake_cache = FakeCache() + + async with FakeAlertmanagerServerWithoutSilence() as fake_alertmanager_server: + port = fake_alertmanager_server.port + alertmanager_client = AlertmanagerClient( + f"http://localhost:{port}", fake_cache + ) + async with aiotools.closing_async(alertmanager_client): + await alertmanager_client.create_silence("fingerprint1", "user1", 86400) + silence_id2 = await alertmanager_client.update_silence( + "fingerprint1", "user2", 864000 + ) + silence2 = await alertmanager_client.get_silence("silence2") + + self.assertEqual("silence2", silence_id2) + self.assertEqual( + { + "id": "silence2", + "status": {"state": "active"}, + "matchers": [ + { + "name": "alertname", + "value": "alert1", + "isRegex": False, + "isEqual": True, + } + ], + "createdBy": "user2", + "startsAt": "1970-01-01T00:00:00", + "endsAt": "1970-01-11T00:00:00", + "comment": "Acknowledge alert from Matrix", + }, + silence2, + ) + self.assertEqual({"fingerprint1": ("silence2", 864000)}, fake_cache.cache) + + @patch.object(matrix_alertbot.alertmanager.AlertmanagerClient, "update_silence") + @patch.object(matrix_alertbot.alertmanager.AlertmanagerClient, "create_silence") + async def test_create_or_update_silence_with_duration_and_silence_not_found( + self, fake_create_silence: Mock, fake_update_silence: Mock + ) -> None: + fake_cache = Mock(spec=Cache) + fake_update_silence.side_effect = update_silence_raise_silence_not_found + fake_create_silence.return_value = "silence1" + + alertmanager_client = AlertmanagerClient("http://localhost", fake_cache) + async with aiotools.closing_async(alertmanager_client): + silence_id1 = await alertmanager_client.create_or_update_silence( + "fingerprint1", "user", 86400 + ) + + self.assertEqual("silence1", silence_id1) + fake_update_silence.assert_called_once_with("fingerprint1", "user", 86400) + fake_create_silence.assert_called_once_with("fingerprint1", "user", 86400) + + @patch.object(matrix_alertbot.alertmanager.AlertmanagerClient, "update_silence") + @patch.object(matrix_alertbot.alertmanager.AlertmanagerClient, "create_silence") + async def test_create_or_update_silence_with_duration_and_silence_found( + self, fake_create_silence: Mock, fake_update_silence: Mock + ) -> None: + fake_cache = Mock(spec=Cache) + fake_update_silence.return_value = "silence1" + + alertmanager_client = AlertmanagerClient("http://localhost", fake_cache) + async with aiotools.closing_async(alertmanager_client): + silence_id1 = await alertmanager_client.create_or_update_silence( + "fingerprint1", "user", 86400 + ) + + self.assertEqual("silence1", silence_id1) + fake_update_silence.assert_called_once_with("fingerprint1", "user", 86400) + fake_create_silence.assert_not_called() + @freeze_time(datetime.utcfromtimestamp(0)) async def test_create_silence_without_duration(self) -> None: fake_cache = Mock(spec=Cache) @@ -525,6 +607,43 @@ class AlertmanagerClientTestCase(unittest.IsolatedAsyncioTestCase): ) self.assertEqual({"fingerprint1": ("silence2", None)}, fake_cache.cache) + @patch.object(matrix_alertbot.alertmanager.AlertmanagerClient, "update_silence") + @patch.object(matrix_alertbot.alertmanager.AlertmanagerClient, "create_silence") + async def test_create_or_update_silence_without_duration_and_silence_not_found( + self, fake_create_silence: Mock, fake_update_silence: Mock + ) -> None: + fake_cache = Mock(spec=Cache) + fake_update_silence.side_effect = update_silence_raise_silence_not_found + fake_create_silence.return_value = "silence1" + + alertmanager_client = AlertmanagerClient("http://localhost", fake_cache) + async with aiotools.closing_async(alertmanager_client): + silence_id1 = await alertmanager_client.create_or_update_silence( + "fingerprint1", "user" + ) + + self.assertEqual("silence1", silence_id1) + fake_update_silence.assert_called_once_with("fingerprint1", "user", None) + fake_create_silence.assert_called_once_with("fingerprint1", "user", None) + + @patch.object(matrix_alertbot.alertmanager.AlertmanagerClient, "update_silence") + @patch.object(matrix_alertbot.alertmanager.AlertmanagerClient, "create_silence") + async def test_create_or_update_silence_without_duration_and_silence_found( + self, fake_create_silence: Mock, fake_update_silence: Mock + ) -> None: + fake_cache = Mock(spec=Cache) + fake_update_silence.return_value = "silence1" + + alertmanager_client = AlertmanagerClient("http://localhost", fake_cache) + async with aiotools.closing_async(alertmanager_client): + silence_id1 = await alertmanager_client.create_or_update_silence( + "fingerprint1", "user" + ) + + self.assertEqual("silence1", silence_id1) + fake_update_silence.assert_called_once_with("fingerprint1", "user", None) + fake_create_silence.assert_not_called() + @freeze_time(datetime.utcfromtimestamp(0)) async def test_create_silence_with_max_duration(self) -> None: fake_cache = Mock(spec=Cache) diff --git a/tests/test_command.py b/tests/test_command.py index d578933..5ee19f3 100644 --- a/tests/test_command.py +++ b/tests/test_command.py @@ -32,7 +32,7 @@ def cache_get_item(key: str) -> str: async def create_silence( - fingerprint: str, user: str, seconds: Optional[int] = None + fingerprint: str, user: str, duration_seconds: Optional[int] = None ) -> str: if fingerprint == "fingerprint1": return "silence1" @@ -42,7 +42,7 @@ async def create_silence( async def create_silence_raise_alertmanager_error( - fingerprint: str, user: str, seconds: Optional[int] = None + fingerprint: str, user: str, duration_seconds: Optional[int] = None ) -> str: if fingerprint == "fingerprint1": raise AlertmanagerError @@ -50,7 +50,7 @@ async def create_silence_raise_alertmanager_error( async def create_silence_raise_alert_not_found_error( - fingerprint: str, user: str, seconds: Optional[int] = None + fingerprint: str, user: str, duration_seconds: Optional[int] = None ) -> str: if fingerprint == "fingerprint1": raise AlertNotFoundError @@ -205,11 +205,10 @@ class CommandTestCase(unittest.IsolatedAsyncioTestCase): async def test_ack_without_duration(self, fake_send_text_to_room: Mock) -> None: """Tests the callback for InviteMemberEvents""" # Tests that the bot attempts to join a room after being invited to it - fake_cache_dict = { - self.fake_alert_event_id: "fingerprint1", - } + fake_cache_dict = {self.fake_alert_event_id: "fingerprint1"} self.fake_cache.__getitem__.side_effect = fake_cache_dict.__getitem__ + self.fake_alertmanager_client.create_or_update_silence.return_value = "silence1" command = AckAlertCommand( self.fake_matrix_client, @@ -224,7 +223,7 @@ class CommandTestCase(unittest.IsolatedAsyncioTestCase): await command.process() # Check that we attempted to create silences - self.fake_alertmanager_client.create_silence.assert_called_once_with( + self.fake_alertmanager_client.create_or_update_silence.assert_called_once_with( "fingerprint1", self.fake_sender, None ) fake_send_text_to_room.assert_called_once_with( @@ -241,9 +240,8 @@ class CommandTestCase(unittest.IsolatedAsyncioTestCase): async def test_ack_with_duration(self, fake_send_text_to_room: Mock) -> None: """Tests the callback for InviteMemberEvents""" # Tests that the bot attempts to join a room after being invited to it - fake_cache_dict = { - self.fake_alert_event_id: "fingerprint1", - } + fake_cache_dict = {self.fake_alert_event_id: "fingerprint1"} + self.fake_alertmanager_client.create_or_update_silence.return_value = "silence1" self.fake_cache.__getitem__.side_effect = fake_cache_dict.__getitem__ @@ -261,7 +259,7 @@ class CommandTestCase(unittest.IsolatedAsyncioTestCase): await command.process() # Check that we attempted to create silences - self.fake_alertmanager_client.create_silence.assert_called_once_with( + self.fake_alertmanager_client.create_or_update_silence.assert_called_once_with( "fingerprint1", self.fake_sender, 864000 ) fake_send_text_to_room.assert_called_once_with( @@ -285,6 +283,9 @@ class CommandTestCase(unittest.IsolatedAsyncioTestCase): } self.fake_cache.__getitem__.side_effect = fake_cache_dict.__getitem__ + self.fake_alertmanager_client.create_or_update_silence.side_effect = ( + create_silence_raise_alertmanager_error + ) command = AckAlertCommand( self.fake_matrix_client, @@ -296,14 +297,10 @@ class CommandTestCase(unittest.IsolatedAsyncioTestCase): self.fake_event_id, self.fake_alert_event_id, ) - - self.fake_alertmanager_client.create_silence.side_effect = ( - create_silence_raise_alertmanager_error - ) await command.process() # Check that we attempted to create silences - self.fake_alertmanager_client.create_silence.assert_called_once_with( + self.fake_alertmanager_client.create_or_update_silence.assert_called_once_with( "fingerprint1", self.fake_sender, None ) fake_send_text_to_room.assert_called_once_with( @@ -325,6 +322,9 @@ class CommandTestCase(unittest.IsolatedAsyncioTestCase): } self.fake_cache.__getitem__.side_effect = fake_cache_dict.__getitem__ + self.fake_alertmanager_client.create_or_update_silence.side_effect = ( + create_silence_raise_alert_not_found_error + ) command = AckAlertCommand( self.fake_matrix_client, @@ -336,14 +336,10 @@ class CommandTestCase(unittest.IsolatedAsyncioTestCase): self.fake_event_id, self.fake_alert_event_id, ) - - self.fake_alertmanager_client.create_silence.side_effect = ( - create_silence_raise_alert_not_found_error - ) await command.process() # Check that we attempted to create silences - self.fake_alertmanager_client.create_silence.assert_called_once_with( + self.fake_alertmanager_client.create_or_update_silence.assert_called_once_with( "fingerprint1", self.fake_sender, None ) fake_send_text_to_room.assert_called_once_with(