From bfc605c20a626fab80aaff766a3f556ab7903f4a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jos=C3=A9=20S=C3=A1nchez-Gallego?= Date: Sat, 21 Dec 2024 19:58:03 -0800 Subject: [PATCH 1/2] Add notifications module --- pyproject.toml | 3 +- src/lvmopstools/data/config.yaml | 14 ++ .../data/critical_error_email_template.html | 94 ++++++++ src/lvmopstools/notifications.py | 204 ++++++++++++++++++ tests/data/test_config.yaml | 20 ++ tests/test_notifications.py | 162 ++++++++++++++ uv.lock | 2 + 7 files changed, 498 insertions(+), 1 deletion(-) create mode 100644 src/lvmopstools/data/critical_error_email_template.html create mode 100644 src/lvmopstools/notifications.py create mode 100644 tests/test_notifications.py diff --git a/pyproject.toml b/pyproject.toml index e595c28..706f403 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -23,6 +23,7 @@ dependencies = [ "aiocache>=0.12.3", "slack-sdk>=3.34.0", "aiohttp>=3.11.11", + "jinja2>=3.1.5", ] [project.optional-dependencies] @@ -92,7 +93,7 @@ typing = ["typing"] sdss = ["sdsstools", "clu", "drift"] [tool.pytest.ini_options] -addopts = "--cov lvmopstools --cov-report xml --cov-report html --cov-report term -W ignore" +addopts = "--cov lvmopstools --cov-report xml --cov-report html --cov-report term -W ignore --capture=tee-sys" asyncio_mode = "auto" asyncio_default_fixture_loop_scope = 'function' diff --git a/src/lvmopstools/data/config.yaml b/src/lvmopstools/data/config.yaml index dfbd3a9..61f9c5e 100644 --- a/src/lvmopstools/data/config.yaml +++ b/src/lvmopstools/data/config.yaml @@ -15,6 +15,20 @@ pubsub: exchange_name: lvmops routing_key: data +notifications: + critical: + email_template: data/critical_error_email_template.html + email_recipients: + - lvm-critical@sdss.org + email_from: LVM Critical Alerts + email_reply_to: lvm-critical@sdss.org + smtp_server: + host: smtp.lco.cl + port: 25 + tls: false + username: null + password: null + devices: thermistors: host: 10.8.38.180 diff --git a/src/lvmopstools/data/critical_error_email_template.html b/src/lvmopstools/data/critical_error_email_template.html new file mode 100644 index 0000000..e55b77d --- /dev/null +++ b/src/lvmopstools/data/critical_error_email_template.html @@ -0,0 +1,94 @@ + + + + + + +
+ + + +
+

+ + SDSS-V LVM Critical Alert + +

+
+
+ +
+ A critical error has been raised in the LVM system. Please see the details below. +
+ +
+ 
{{- message }}
+ + + + + diff --git a/src/lvmopstools/notifications.py b/src/lvmopstools/notifications.py new file mode 100644 index 0000000..7f99a7f --- /dev/null +++ b/src/lvmopstools/notifications.py @@ -0,0 +1,204 @@ +#!/usr/bin/env python +# -*- coding: utf-8 -*- +# +# @Author: José Sánchez-Gallego (gallegoj@uw.edu) +# @Date: 2024-12-21 +# @Filename: notifications.py +# @License: BSD 3-clause (http://www.opensource.org/licenses/BSD-3-Clause) + +from __future__ import annotations + +import enum +import pathlib +import smtplib +import sys +from email.mime.multipart import MIMEMultipart +from email.mime.text import MIMEText + +from typing import Any, Sequence, cast + +from jinja2 import Environment, FileSystemLoader + +from lvmopstools import config +from lvmopstools.slack import post_message as post_to_slack + + +__all__ = ["send_notification", "send_critical_error_email", "NotificationLevel"] + + +class NotificationLevel(enum.Enum): + """Allowed notification levels.""" + + DEBUG = "DEBUG" + INFO = "INFO" + WARNING = "WARNING" + ERROR = "ERROR" + CRITICAL = "CRITICAL" + + +async def send_notification( + message: str, + level: NotificationLevel | str = NotificationLevel.INFO, + slack: bool = True, + slack_channels: str | Sequence[str] | None = None, + email_on_critical: bool = True, + slack_extra_params: dict[str, Any] = {}, + email_params: dict[str, Any] = {}, +): + """Creates a new notification. + + Parameters + ---------- + message + The message of the notification. Can be formatted in Markdown. + level + The level of the notification. + slack + Whether to send the notification to Slack. + slack_channels + The Slack channel where to send the notification. If not provided, the default + channel is used. Can be set to false to disable sending the Slack notification. + email_on_critical + Whether to send an email if the notification level is ``CRITICAL``. + slack_extra_params + A dictionary of extra parameters to pass to ``post_message``. + email_params + A dictionary of extra parameters to pass to :obj:`.send_critical_error_email`. + + Returns + ------- + message + The message that was sent. + + """ + + if isinstance(level, str): + level = NotificationLevel(level.upper()) + else: + level = NotificationLevel(level) + + send_email = email_on_critical and level == NotificationLevel.CRITICAL + + if send_email: + try: + send_critical_error_email(message, **email_params) + except Exception as ee: + print(f"Error sending critical error email: {ee}", file=sys.stderr) + + if slack: + channels: set[str] = set() + + if isinstance(slack_channels, str): + channels.add(slack_channels) + elif isinstance(slack_channels, Sequence): + channels.update(slack_channels) + else: + channels.add(config["slack.default_channel"]) + + # We send the message to the default channel plus any other channel that + # matches the level of the notification. + level_channels = cast(dict[str, str], config["slack.level_channels"]) + if level.value in level_channels: + channels.add(level_channels[level.value]) + + # Send Slack message(s) + for channel in channels: + mentions = ( + ["@channel"] + if level == NotificationLevel.CRITICAL + or level == NotificationLevel.ERROR + else [] + ) + try: + await post_to_slack( + message, + channel=channel, + mentions=mentions, + **slack_extra_params, + ) + except Exception as se: + print(f"Error sending Slack message: {se}", file=sys.stderr) + + return message + + +def send_critical_error_email( + message: str, + host: str | None = None, + port: int | None = None, + tls: bool | None = None, + username: str | None = None, + password: str | None = None, +): + """Sends a critical error email. + + Parameters + ---------- + message + The message to send. + host + The SMTP server host. + port + The SMTP server port. + tls + Whether to use TLS for authentication. + username + The SMTP server username. + password + The SMTP server password. + + """ + + root = pathlib.Path(__file__).parent + template = root / config["notifications.critical.email_template"] + loader = FileSystemLoader(template.parent) + + env = Environment( + loader=loader, + lstrip_blocks=True, + trim_blocks=True, + ) + html_template = env.get_template(template.name) + + html_message = html_template.render(message=message.strip()) + + recipients = config["notifications.critical.email_recipients"] + from_address = config["notifications.critical.email_from"] + + email_reply_to = config["notifications.critical.email_reply_to"] + + msg = MIMEMultipart("alternative" if html_message else "mixed") + msg["Subject"] = "LVM Critical Alert" + msg["From"] = from_address + msg["To"] = ", ".join(recipients) + msg["Reply-To"] = email_reply_to + + plaintext_email = f"""A critical alert was raised in the LVM system. + +The error message is shown below: + +{message} + +""" + msg.attach(MIMEText(plaintext_email, "plain")) + + html = MIMEText(html_message, "html") + msg.attach(html) + + smtp_host = host or config["notifications.smtp_server.host"] + smtp_port = port or config["notifications.smtp_server.port"] + smpt_tls = tls if tls is not None else config["notifications.smtp_server.tls"] + smtp_username = username or config["notifications.smtp_server.username"] + smtp_password = password or config["notifications.smtp_server.password"] + + with smtplib.SMTP(host=smtp_host, port=smtp_port) as smtp: + if smpt_tls is True or (smpt_tls is None and smtp_port == 587): + # See https://gist.github.com/jamescalam/93d915e4de12e7f09834ae73bdf37299 + smtp.ehlo() + smtp.starttls() + + if smtp_password is not None and smtp_password is not None: + smtp.login(smtp_username, smtp_password) + else: + raise ValueError("username and password must be provided for TLS.") + smtp.sendmail(from_address, recipients, msg.as_string()) diff --git a/tests/data/test_config.yaml b/tests/data/test_config.yaml index 5631fb2..7af6b42 100644 --- a/tests/data/test_config.yaml +++ b/tests/data/test_config.yaml @@ -1,3 +1,23 @@ +slack: + token: null + default_channel: lvm-notifications + level_channels: + CRITICAL: lvm-alerts + +notifications: + critical: + email_template: data/critical_error_email_template.html + email_recipients: + - lvm-critical@sdss.org + email_from: LVM Critical Alerts + email_reply_to: lvm-critical@sdss.org + smtp_server: + host: smtp.lco.cl + port: 25 + tls: false + username: null + password: null + pubsub: connection_string: amqp://guest:guest@localhost: exchange_name: lvmops diff --git a/tests/test_notifications.py b/tests/test_notifications.py new file mode 100644 index 0000000..98a65d4 --- /dev/null +++ b/tests/test_notifications.py @@ -0,0 +1,162 @@ +#!/usr/bin/env python +# -*- coding: utf-8 -*- +# +# @Author: José Sánchez-Gallego (gallegoj@uw.edu) +# @Date: 2024-12-21 +# @Filename: test_notifications.py +# @License: BSD 3-clause (http://www.opensource.org/licenses/BSD-3-Clause) + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import pytest + +from lvmopstools.notifications import ( + NotificationLevel, + send_critical_error_email, + send_notification, +) + + +if TYPE_CHECKING: + from pytest_mock import MockerFixture + + +@pytest.mark.parametrize("level", ["INFO", "CRITICAL", NotificationLevel.DEBUG]) +async def test_send_notification(mocker: MockerFixture, level: str | NotificationLevel): + slack_mock = mocker.patch("lvmopstools.notifications.post_to_slack") + email_mock = mocker.patch("lvmopstools.notifications.send_critical_error_email") + + message = await send_notification("test message", level=level) + + assert message == "test message" + slack_mock.assert_called_with( + "test message", + channel=mocker.ANY, + mentions=["@channel"] if level == "CRITICAL" else [], + ) + + if level == "CRITICAL": + email_mock.assert_called() + else: + email_mock.assert_not_called() + + +@pytest.mark.parametrize( + "level,channels,n_calls,", + [ + ("INFO", None, 1), + ("INFO", "test-channel", 1), + ("INFO", ["test-channel", "test-channel2"], 2), + ("CRITICAL", None, 2), + ("CRITICAL", "test-channel", 2), + ], +) +async def test_send_notification_channels( + mocker: MockerFixture, + level: str | NotificationLevel, + channels: str | list[str], + n_calls: int, +): + slack_mock = mocker.patch("lvmopstools.notifications.post_to_slack") + email_mock = mocker.patch("lvmopstools.notifications.send_critical_error_email") + + await send_notification( + "test message", + level=level, + slack_channels=channels, + slack=True, + ) + + assert slack_mock.call_count == n_calls + + if level == "CRITICAL": + email_mock.assert_called() + else: + email_mock.assert_not_called() + + +async def test_send_notification_no_slack(mocker: MockerFixture): + slack_mock = mocker.patch("lvmopstools.notifications.post_to_slack") + + await send_notification("test message", slack=False) + + slack_mock.assert_not_called() + + +async def test_send_notification_no_email(mocker: MockerFixture): + email_mock = mocker.patch("lvmopstools.notifications.send_critical_error_email") + + await send_notification( + "test message", + email_on_critical=False, + level="CRITICAL", + slack=False, + ) + + email_mock.assert_not_called() + + +async def test_post_to_slack_fails( + mocker: MockerFixture, + capsys: pytest.CaptureFixture, +): + mocker.patch("lvmopstools.notifications.post_to_slack", side_effect=ValueError()) + + await send_notification("test message", slack=True) + + stderr = capsys.readouterr().err + assert "Error sending Slack message" in stderr + + +async def test_send_email_fails( + mocker: MockerFixture, + capsys: pytest.CaptureFixture, +): + mocker.patch( + "lvmopstools.notifications.send_critical_error_email", + side_effect=ValueError(), + ) + + await send_notification("test message", level="CRITICAL") + + stderr = capsys.readouterr().err + assert "Error sending critical error email" in stderr + + +def test_send_email(mocker: MockerFixture): + smtp_mock = mocker.patch("lvmopstools.notifications.smtplib.SMTP", autospec=True) + sendmail_mock = smtp_mock.return_value.__enter__.return_value.sendmail + + send_critical_error_email("test message") + + sendmail_mock.assert_called() + assert "Content-Type: multipart/alternative" in sendmail_mock.call_args[0][-1] + assert "" in sendmail_mock.call_args[0][-1] + + +def test_send_email_tls(mocker: MockerFixture): + smtp_mock = mocker.patch("lvmopstools.notifications.smtplib.SMTP", autospec=True) + sendmail_mock = smtp_mock.return_value.__enter__.return_value.sendmail + + send_critical_error_email( + "test message", + tls=True, + username="test", + password="test", + ) + + sendmail_mock.assert_called() + smtp_mock.return_value.__enter__.return_value.starttls.assert_called() + + +def test_send_email_tls_no_password(mocker: MockerFixture): + mocker.patch("lvmopstools.notifications.smtplib.SMTP", autospec=True) + + with pytest.raises(ValueError): + send_critical_error_email( + "test message", + tls=True, + username="test", + ) diff --git a/uv.lock b/uv.lock index 9292e6c..6c7cbe9 100644 --- a/uv.lock +++ b/uv.lock @@ -958,6 +958,7 @@ dependencies = [ { name = "aiohttp" }, { name = "asyncudp" }, { name = "httpx" }, + { name = "jinja2" }, { name = "polars" }, { name = "pydantic" }, { name = "sdss-clu" }, @@ -1023,6 +1024,7 @@ requires-dist = [ { name = "cachetools", marker = "extra == 'schedule'", specifier = ">=5.5.0" }, { name = "httpx", specifier = ">=0.27.2" }, { name = "influxdb-client", marker = "extra == 'influxdb'", specifier = ">=1.47.0" }, + { name = "jinja2", specifier = ">=3.1.5" }, { name = "kubernetes", marker = "extra == 'kubernetes'", specifier = ">=31.0.0" }, { name = "polars", specifier = ">=1.13.0" }, { name = "pydantic", specifier = ">=2.10.3" }, From 75234ef3b0e6357b3b8a31dceca2cbaae2a0f574 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jos=C3=A9=20S=C3=A1nchez-Gallego?= Date: Sat, 21 Dec 2024 20:02:23 -0800 Subject: [PATCH 2/2] Update changelog and docs --- CHANGELOG.md | 1 + docs/sphinx/api.rst | 6 ++++++ 2 files changed, 7 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 5a1651c..ba664af 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,6 +10,7 @@ * [#10](https://vscode.dev/github/sdss/lvmopstools/pull/10) Added a `pubsub` module with tools to emit and subscribe to events using RabbitMQ. * [#11](https://vscode.dev/github/sdss/lvmopstools/pull/11) Added a `slack` module with tools to send messages to Slack. +* [#12](https://vscode.dev/github/sdss/lvmopstools/pull/12) Added a `notifications` module. * Added `ephemeris.is_sun_up`. diff --git a/docs/sphinx/api.rst b/docs/sphinx/api.rst index 5ace18b..938047c 100644 --- a/docs/sphinx/api.rst +++ b/docs/sphinx/api.rst @@ -20,6 +20,12 @@ InfluxDB .. autofunction:: lvmopstools.influxdb.query_influxdb +Notifications +------------- + +.. automodule:: lvmopstools.notifications + :members: + PubSub ------