dify-1.8.1/api/tests/unit_tests/services/test_alert_service.py

"""Unit tests for alert service."""

import json
from unittest.mock import MagicMock, patch

import pytest

from services.alert_service import AlertService


class TestAlertService:
    """Test cases for AlertService."""

    def test_check_threshold_greater_than(self):
        """Test check_threshold with > operator."""
        assert AlertService.check_threshold(10.0, ">", 5.0) is True
        assert AlertService.check_threshold(5.0, ">", 5.0) is False
        assert AlertService.check_threshold(3.0, ">", 5.0) is False
        assert AlertService.check_threshold(None, ">", 5.0) is False

    def test_check_threshold_less_than(self):
        """Test check_threshold with < operator."""
        assert AlertService.check_threshold(3.0, "<", 5.0) is True
        assert AlertService.check_threshold(5.0, "<", 5.0) is False
        assert AlertService.check_threshold(10.0, "<", 5.0) is False
        assert AlertService.check_threshold(None, "<", 5.0) is False

    def test_check_threshold_greater_than_or_equal(self):
        """Test check_threshold with >= operator."""
        assert AlertService.check_threshold(10.0, ">=", 5.0) is True
        assert AlertService.check_threshold(5.0, ">=", 5.0) is True
        assert AlertService.check_threshold(3.0, ">=", 5.0) is False

    def test_check_threshold_less_than_or_equal(self):
        """Test check_threshold with <= operator."""
        assert AlertService.check_threshold(3.0, "<=", 5.0) is True
        assert AlertService.check_threshold(5.0, "<=", 5.0) is True
        assert AlertService.check_threshold(10.0, "<=", 5.0) is False

    def test_check_threshold_equal(self):
        """Test check_threshold with = operator."""
        assert AlertService.check_threshold(5.0, "=", 5.0) is True
        assert AlertService.check_threshold(5.1, "=", 5.0) is False
        assert AlertService.check_threshold(4.9, "=", 5.0) is False

    def test_check_threshold_unknown_operator(self):
        """Test check_threshold with unknown operator."""
        assert AlertService.check_threshold(10.0, "unknown", 5.0) is False

    @patch("services.alert_service.redis_client")
    def test_create_alert_rule(self, mock_redis):
        """Test creating an alert rule."""
        # Mock app query
        with patch("services.alert_service.db") as mock_db:
            mock_app = MagicMock()
            mock_app.id = "app-123"
            mock_app.tenant_id = "tenant-456"
            mock_db.session.query.return_value.filter.return_value.first.return_value = mock_app

            # Mock redis operations
            mock_redis.hset.return_value = True
            mock_redis.expire.return_value = True
            mock_redis.sadd.return_value = True
            mock_redis.hgetall.return_value = {}

            data = {
                "name": "Test Alert",
                "description": "Test description",
                "metric_type": "error_rate",
                "operator": ">",
                "threshold_value": 5.0,
                "notification_channels": {"email": ["test@example.com"]},
            }

            result = AlertService.create_alert_rule("app-123", data, "user-789")

            # Verify redis operations were called
            assert mock_redis.hset.called
            assert mock_redis.expire.called
            assert mock_redis.sadd.called

    @patch("services.alert_service.redis_client")
    def test_get_alert_rules(self, mock_redis):
        """Test getting alert rules for an app."""
        # Mock redis operations
        mock_redis.smembers.return_value = [b"rule-1", b"rule-2"]
        mock_redis.hgetall.side_effect = [
            {
                b"id": b"rule-1",
                b"name": b"Rule 1",
                b"metric_type": b"error_rate",
                b"operator": b">",
                b"threshold_value": b"5.0",
                b"is_enabled": b"true",
                b"notification_channels": b'{"email": []}',
                b"check_interval_minutes": b"5",
                b"consecutive_triggers": b"2",
                b"trigger_count": b"0",
                b"consecutive_trigger_count": b"0",
            },
            {
                b"id": b"rule-2",
                b"name": b"Rule 2",
                b"metric_type": b"avg_response_time",
                b"operator": b">",
                b"threshold_value": b"1000.0",
                b"is_enabled": b"false",
                b"notification_channels": b'{"email": []}',
                b"check_interval_minutes": b"10",
                b"consecutive_triggers": b"3",
                b"trigger_count": b"5",
                b"consecutive_trigger_count": b"1",
            },
        ]

        rules = AlertService.get_alert_rules("app-123")

        assert len(rules) == 2
        assert rules[0]["id"] == "rule-1"
        assert rules[0]["is_enabled"] is True
        assert rules[0]["threshold_value"] == 5.0
        assert rules[1]["id"] == "rule-2"
        assert rules[1]["is_enabled"] is False
        assert rules[1]["trigger_count"] == 5

    @patch("services.alert_service.redis_client")
    @patch("services.alert_service.NotificationService")
    @patch("services.alert_service.db")
    def test_trigger_alert(self, mock_db, mock_notification_service, mock_redis):
        """Test triggering an alert."""
        # Mock app query
        mock_app = MagicMock()
        mock_app.name = "Test App"
        mock_db.session.query.return_value.filter.return_value.first.return_value = mock_app

        # Mock notification service
        mock_notification_service.send_alert_notification.return_value = {
            "email": {"status": "success"},
        }

        rule = {
            "id": "rule-123",
            "app_id": "app-123",
            "metric_type": "error_rate",
            "operator": ">",
            "threshold_value": 5.0,
            "notification_channels": {"email": ["test@example.com"]},
            "trigger_count": 0,
        }

        AlertService.trigger_alert(rule, 10.0)

        # Verify notification was sent
        mock_notification_service.send_alert_notification.assert_called_once()

        # Verify log was created
        assert mock_redis.hset.called
        assert mock_redis.zadd.called

    @patch("services.alert_service.redis_client")
    def test_get_alert_logs(self, mock_redis):
        """Test getting alert logs for a rule."""
        # Mock redis operations
        mock_redis.zrevrange.return_value = [b"log-1", b"log-2"]
        mock_redis.hgetall.side_effect = [
            {
                b"id": b"log-1",
                b"alert_rule_id": b"rule-123",
                b"metric_value": b"10.5",
                b"threshold_value": b"5.0",
                b"notification_status": b'{"email": {"status": "success"}}',
                b"error_message": b"",
                b"triggered_at": b"2024-10-28T10:00:00Z",
            },
            {
                b"id": b"log-2",
                b"alert_rule_id": b"rule-123",
                b"metric_value": b"8.2",
                b"threshold_value": b"5.0",
                b"notification_status": b'{"email": {"status": "success"}}',
                b"error_message": b"",
                b"triggered_at": b"2024-10-28T09:00:00Z",
            },
        ]

        logs = AlertService.get_alert_logs("rule-123", limit=100)

        assert len(logs) == 2
        assert logs[0]["id"] == "log-1"
        assert logs[0]["metric_value"] == 10.5
        assert logs[1]["id"] == "log-2"
        assert logs[1]["metric_value"] == 8.2