matrix-alertbot/tests/test_webhook.py

import unittest
from typing import Dict
from unittest.mock import Mock, patch

import aiohttp.test_utils
import nio
from aiohttp import web
from diskcache import Cache
from nio import LocalProtocolError, RoomSendResponse

import matrix_alertbot.webhook
from matrix_alertbot.config import Config
from matrix_alertbot.webhook import Webhook


def send_text_to_room_raise_error(
    client: nio.AsyncClient, room_id: str, plaintext: str, html: str, notice: bool
) -> RoomSendResponse:
    raise LocalProtocolError()


class WebhookApplicationTestCase(aiohttp.test_utils.AioHTTPTestCase):
    async def get_application(self) -> web.Application:
        self.fake_client = Mock(spec=nio.AsyncClient)
        self.fake_cache = Mock(spec=Cache)

        self.fake_config = Mock(spec=Config)
        self.fake_config.port = aiohttp.test_utils.unused_port()
        self.fake_config.address = "localhost"
        self.fake_config.socket = "webhook.sock"
        self.fake_config.room_id = "!abcdefg:example.com"
        self.fake_config.cache_expire_time = 0

        self.fake_alerts = {
            "alerts": [
                {
                    "fingerprint": "fingerprint1",
                    "generatorURL": "http://example.com/alert1",
                    "status": "firing",
                    "labels": {
                        "alertname": "alert1",
                        "severity": "critical",
                        "job": "job1",
                    },
                    "annotations": {"description": "some description1"},
                },
                {
                    "fingerprint": "fingerprint2",
                    "generatorURL": "http://example.com/alert2",
                    "status": "resolved",
                    "labels": {
                        "alertname": "alert2",
                        "severity": "warning",
                        "job": "job2",
                    },
                    "annotations": {"description": "some description2"},
                },
            ]
        }

        webhook = Webhook(self.fake_client, self.fake_cache, self.fake_config)
        return webhook.app

    @patch.object(matrix_alertbot.webhook, "send_text_to_room")
    async def test_post_alert(self, fake_send_text_to_room: Mock) -> None:
        data = self.fake_alerts
        async with self.client.request("POST", "/alert", json=data) as response:
            self.assertEqual(200, response.status)
        fake_send_text_to_room.assert_called_once_with(
            self.fake_client,
            self.fake_config.room_id,
            "[🔥 CRITICAL] alert1: some description1\n"
            "[🥦 RESOLVED] alert2: some description2",
            "<font color='#dc3545'><b>[🔥 CRITICAL]</b></font> "
            "<a href='http://example.com/alert1'>alert1</a> (job1)<br/>"
            "some description1<br/>\n"
            "<font color='#33cc33'><b>[🥦 RESOLVED]</b></font> "
            "<a href='http://example.com/alert2'>alert2</a> (job2)<br/>"
            "some description2",
            notice=False,
        )

    @patch.object(matrix_alertbot.webhook, "send_text_to_room")
    async def test_post_alert_with_empty_data(
        self, fake_send_text_to_room: Mock
    ) -> None:
        async with self.client.request("POST", "/alert", json={}) as response:
            self.assertEqual(400, response.status)
            error_msg = await response.text()
            self.assertEqual("Data must contain 'alerts' key.", error_msg)
        fake_send_text_to_room.assert_not_called()

    @patch.object(matrix_alertbot.webhook, "send_text_to_room")
    async def test_post_alert_with_empty_alerts(
        self, fake_send_text_to_room: Mock
    ) -> None:
        data: Dict = {"alerts": []}
        async with self.client.request("POST", "/alert", json=data) as response:
            self.assertEqual(400, response.status)
            error_msg = await response.text()
            self.assertEqual("Alerts cannot be empty.", error_msg)
        fake_send_text_to_room.assert_not_called()

    @patch.object(matrix_alertbot.webhook, "send_text_to_room")
    async def test_post_alert_with_invalid_alerts(
        self, fake_send_text_to_room: Mock
    ) -> None:
        data = {"alerts": "invalid"}
        async with self.client.request("POST", "/alert", json=data) as response:
            self.assertEqual(400, response.status)
            error_msg = await response.text()
            self.assertEqual("Alerts must be a list.", error_msg)
        fake_send_text_to_room.assert_not_called()

    @patch.object(matrix_alertbot.webhook, "send_text_to_room")
    async def test_post_alert_with_empty_items(
        self, fake_send_text_to_room: Mock
    ) -> None:
        data: Dict = {"alerts": [{}]}
        async with self.client.request("POST", "/alert", json=data) as response:
            self.assertEqual(400, response.status)
            error_msg = await response.text()
            self.assertEqual("Invalid alert: {}.", error_msg)
        fake_send_text_to_room.assert_not_called()

    @patch.object(
        matrix_alertbot.webhook,
        "send_text_to_room",
        side_effect=send_text_to_room_raise_error,
    )
    async def test_post_alert_with_send_error(
        self, fake_send_text_to_room: Mock
    ) -> None:
        data = self.fake_alerts
        async with self.client.request("POST", "/alert", json=data) as response:
            self.assertEqual(500, response.status)
            error_msg = await response.text()
            self.assertEqual(
                "An error occured when sending alerts to Matrix room.", error_msg
            )
        fake_send_text_to_room.assert_called_once()


class WebhookServerTestCase(unittest.IsolatedAsyncioTestCase):
    async def asyncSetUp(self) -> None:
        self.fake_client = Mock(spec=nio.AsyncClient)
        self.fake_cache = Mock(spec=Cache)

        self.fake_config = Mock(spec=Config)
        self.fake_config.port = aiohttp.test_utils.unused_port()
        self.fake_config.address = "localhost"
        self.fake_config.socket = "webhook.sock"
        self.fake_config.room_id = "!abcdefg:example.com"
        self.fake_config.cache_expire_time = 0

    @patch.object(matrix_alertbot.webhook.web, "TCPSite", autospec=True)
    async def test_webhook_start_address_port(self, fake_tcp_site: Mock) -> None:
        webhook = Webhook(self.fake_client, self.fake_cache, self.fake_config)
        await webhook.start()

        fake_tcp_site.assert_called_once_with(
            webhook.runner, self.fake_config.address, self.fake_config.port
        )

        await webhook.close()

    @patch.object(matrix_alertbot.webhook.web, "UnixSite", autospec=True)
    async def test_webhook_start_unix_socket(self, fake_unix_site: Mock) -> None:
        self.fake_config.address = None
        self.fake_config.port = None

        webhook = Webhook(self.fake_client, self.fake_cache, self.fake_config)
        await webhook.start()

        fake_unix_site.assert_called_once_with(webhook.runner, self.fake_config.socket)

        await webhook.close()