Daily Email Limits - Add limits to redis (#1934)

* first commit for email limits?

* fixes

* remove

* fix types

app/dao/services_dao.py

@@ -447,6 +447,25 @@ def fetch_todays_total_sms_count(service_id):
     return 0 if result is None or result.sum_billable_units is None else result.sum_billable_units
 
 
+def fetch_service_email_limit(service_id: uuid.UUID) -> int:
+    return Service.query.get(service_id).message_limit
+
+
+def fetch_todays_total_email_count(service_id: uuid.UUID) -> int:
+    midnight = get_midnight(datetime.now(tz=pytz.utc))
+    result = (
+        db.session.query(func.count(Notification).label("total_email_notifications"))
+        .filter(
+            Notification.service_id == service_id,
+            Notification.key_type != KEY_TYPE_TEST,
+            Notification.created_at > midnight,
+            Notification.notification_type == "email",
+        )
+        .first()
+    )
+    return 0 if result is None else result.total_email_notifications
+
+
 def _stats_for_service_query(service_id):
     return (
         db.session.query(

app/email_limit_utils.py

@@ -0,0 +1,29 @@
+from datetime import timedelta
+from uuid import UUID
+
+from flask import current_app
+from notifications_utils.clients.redis import email_daily_count_cache_key
+
+from app import redis_store
+from app.dao.services_dao import fetch_todays_total_email_count
+
+
+def fetch_todays_email_count(service_id: UUID) -> int:
+    if not current_app.config["REDIS_ENABLED"]:
+        return fetch_todays_total_email_count(service_id)
+
+    cache_key = email_daily_count_cache_key(service_id)
+    total_email_count = redis_store.get(cache_key)
+    if total_email_count is None:
+        total_email_count = fetch_todays_total_email_count(service_id)
+        redis_store.set(cache_key, total_email_count, ex=int(timedelta(hours=2).total_seconds()))
+    return int(total_email_count)
+
+
+def increment_todays_email_count(service_id: UUID, increment_by: int) -> None:
+    if not current_app.config["REDIS_ENABLED"]:
+        return
+
+    fetch_todays_email_count(service_id)  # to make sure it's set in redis
+    cache_key = email_daily_count_cache_key(service_id)
+    redis_store.incrby(cache_key, increment_by)

tests/app/dao/test_services_dao.py

@@ -37,6 +37,7 @@
     dao_suspend_service,
     dao_update_service,
     delete_service_and_all_associated_db_objects,
+    fetch_service_email_limit,
     fetch_todays_total_message_count,
     fetch_todays_total_sms_count,
     get_services_by_partial_name,
@@ -1418,3 +1419,46 @@ def create_email_sms_letter_template():
     template_two = create_template(service=service, template_name="2", template_type="sms")
     template_three = create_template(service=service, template_name="3", template_type="letter")
     return template_one, template_three, template_two
+
+
+class TestServiceEmailLimits:
+    def test_get_email_count_for_service(self, notify_db_session):
+        active_user_1 = create_user(email="[email protected]", state="active")
+        service = Service(
+            name="service_name",
+            email_from="email_from",
+            message_limit=1000,
+            sms_daily_limit=1000,
+            restricted=False,
+            created_by=active_user_1,
+        )
+        dao_create_service(
+            service,
+            active_user_1,
+            service_permissions=[
+                SMS_TYPE,
+                EMAIL_TYPE,
+                INTERNATIONAL_SMS_TYPE,
+            ],
+        )
+        assert fetch_service_email_limit(service.id) == 1000
+
+    def test_dao_fetch_todays_total_message_count_returns_count_for_today(self):
+        service = create_service()
+        email_template = create_template(service=service, template_type="email")
+        save_notification(create_notification(template=email_template, status="created"))
+        assert fetch_todays_total_message_count(service.id) == 1
+
+    def test_dao_fetch_todays_total_message_count_returns_0_when_no_messages_for_today(self):
+        assert fetch_todays_total_message_count(uuid.uuid4()) == 0
+
+    def test_dao_fetch_todays_total_message_count_returns_0_with_yesterday_messages(self):
+        today = datetime.utcnow().date()
+        yesterday = today - timedelta(days=1)
+        notification = save_notification(
+            create_notification(
+                created_at=yesterday,
+                template=create_template(service=create_service(service_name="tester"), template_type="email"),
+            )
+        )
+        assert fetch_todays_total_message_count(notification.service.id) == 0

tests/app/test_email_limit_utils.py

@@ -0,0 +1,46 @@
+import pytest
+from notifications_utils.clients.redis import email_daily_count_cache_key
+
+from app.email_limit_utils import fetch_todays_email_count, increment_todays_email_count
+from tests.conftest import set_config
+
+
+class TestEmailLimits:
+    @pytest.mark.parametrize("redis_value, db_value, expected_result", [(None, 5, 5), ("3", 5, 3)])
+    def test_fetch_todays_requested_email_count(self, client, mocker, sample_service, redis_value, db_value, expected_result):
+        cache_key = email_daily_count_cache_key(sample_service.id)
+        mocker.patch("app.redis_store.get", lambda x: redis_value if x == cache_key else None)
+        mocked_set = mocker.patch("app.redis_store.set")
+        mocker.patch("app.email_limit_utils.fetch_todays_total_email_count", return_value=db_value)
+        # mocker.patch("app.dao.users_dao.user_can_be_archived", return_value=False)
+
+        with set_config(client.application, "REDIS_ENABLED", True):
+            actual_result = fetch_todays_email_count(sample_service.id)
+
+            assert actual_result == expected_result
+            if redis_value is None:
+                assert mocked_set.called_once_with(
+                    cache_key,
+                    db_value,
+                )
+            else:
+                mocked_set.assert_not_called()
+
+    @pytest.mark.parametrize("redis_value, db_value, increment_by", [(None, 5, 5), ("3", 5, 3)])
+    def test_increment_todays_requested_email_count(self, mocker, sample_service, redis_value, db_value, increment_by):
+        cache_key = email_daily_count_cache_key(sample_service.id)
+        mocker.patch("app.redis_store.get", lambda x: redis_value if x == cache_key else None)
+        mocked_set = mocker.patch("app.redis_store.set")
+        mocked_incrby = mocker.patch("app.redis_store.incrby")
+        mocker.patch("app.email_limit_utils.fetch_todays_email_count", return_value=db_value)
+
+        increment_todays_email_count(sample_service.id, increment_by)
+
+        assert mocked_incrby.called_once_with(cache_key, increment_by)
+        if redis_value is None:
+            assert mocked_set.called_once_with(
+                cache_key,
+                db_value,
+            )
+        else:
+            mocked_set.assert_not_called()