Fix daily limit counting when scheduling jobs

.devcontainer/scripts/notify-dev-entrypoint.sh

@@ -34,8 +34,10 @@ cd /workspace
 echo -e "fpath+=/.zfunc" >> ~/.zshrc
 echo -e "autoload -Uz compinit && compinit"
 
-pip install poetry==${POETRY_VERSION} \
-  && poetry --version
+pip install poetry==${POETRY_VERSION}
+export PATH=$PATH:/home/vscode/.local/bin/
+which poetry
+poetry --version
 
 # Initialize poetry autocompletions
 mkdir ~/.zfunc

app/dao/services_dao.py

@@ -23,6 +23,7 @@
     CROWN_ORGANISATION_TYPES,
     EMAIL_TYPE,
     INTERNATIONAL_SMS_TYPE,
+    JOB_STATUS_SCHEDULED,
     KEY_TYPE_TEST,
     NHS_ORGANISATION_TYPES,
     NON_CROWN_ORGANISATION_TYPES,
@@ -424,20 +425,24 @@ def dao_fetch_todays_stats_for_service(service_id):
 
 def fetch_todays_total_message_count(service_id):
     midnight = get_midnight(datetime.now(tz=pytz.utc))
+    scheduled = (
+        db.session.query(func.coalesce(func.sum(Job.notification_count), 0).label("count")).filter(
+            Job.service_id == service_id,
+            Job.job_status == JOB_STATUS_SCHEDULED,
+            Job.scheduled_for >= midnight,
+            Job.scheduled_for < midnight + timedelta(days=1),
+        )
+    ).first()
+
     result = (
-        db.session.query(func.count(Notification.id).label("count"))
-        .filter(
+        db.session.query(func.coalesce(func.count(Notification.id), 0).label("count")).filter(
             Notification.service_id == service_id,
             Notification.key_type != KEY_TYPE_TEST,
-            Notification.created_at > midnight,
+            Notification.created_at >= midnight,
         )
-        .group_by(
-            Notification.notification_type,
-            Notification.status,
-        )
-        .first()
-    )
-    return 0 if result is None else result.count
+    ).first()
+
+    return result.count + scheduled.count
 
 
 def fetch_todays_total_sms_count(service_id):
@@ -461,17 +466,25 @@ def fetch_service_email_limit(service_id: uuid.UUID) -> int:
 
 def fetch_todays_total_email_count(service_id: uuid.UUID) -> int:
     midnight = get_midnight(datetime.now(tz=pytz.utc))
+    scheduled = (
+        db.session.query(func.coalesce(func.sum(Job.notification_count), 0).label("total_scheduled_notifications")).filter(
+            Job.service_id == service_id,
+            Job.job_status == JOB_STATUS_SCHEDULED,
+            Job.scheduled_for > midnight,
+            Job.scheduled_for < midnight + timedelta(hours=23, minutes=59, seconds=59),
+        )
+    ).first()
+
     result = (
-        db.session.query(func.count(Notification.id).label("total_email_notifications"))
-        .filter(
+        db.session.query(func.coalesce(func.count(Notification.id), 0).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
+    ).first()
+
+    return result.total_email_notifications + scheduled.total_scheduled_notifications
 
 
 def _stats_for_service_query(service_id):

app/email_limit_utils.py

@@ -3,6 +3,7 @@
 
 from flask import current_app
 from notifications_utils.clients.redis import email_daily_count_cache_key
+from notifications_utils.decorators import requires_feature
 
 from app import redis_store
 from app.dao.services_dao import fetch_todays_total_email_count
@@ -20,9 +21,15 @@ def fetch_todays_email_count(service_id: UUID) -> int:
     return int(total_email_count)
 
 
+@requires_feature("REDIS_ENABLED")
 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)
+
+
+@requires_feature("REDIS_ENABLED")
+def decrement_todays_email_count(service_id: UUID, decrement_by: int) -> None:
+    fetch_todays_email_count(service_id)
+    cache_key = email_daily_count_cache_key(service_id)
+    redis_store.decrby(cache_key, decrement_by)

app/job/rest.py

@@ -1,3 +1,5 @@
+from datetime import datetime
+
 import dateutil
 from flask import Blueprint, current_app, jsonify, request
 from notifications_utils.recipients import RecipientCSV
@@ -20,6 +22,7 @@
 from app.dao.notifications_dao import get_notifications_for_job
 from app.dao.services_dao import dao_fetch_service_by_id
 from app.dao.templates_dao import dao_get_template_by_id
+from app.email_limit_utils import decrement_todays_email_count
 from app.errors import InvalidRequest, register_errors
 from app.models import (
     EMAIL_TYPE,
@@ -67,7 +70,7 @@ def cancel_job(service_id, job_id):
     job = dao_get_future_scheduled_job_by_id_and_service_id(job_id, service_id)
     job.job_status = JOB_STATUS_CANCELLED
     dao_update_job(job)
-
+    decrement_todays_email_count(service_id, job.notification_count)
     return get_job_by_service_and_job_id(service_id, job_id)
 
 
@@ -137,15 +140,23 @@ def create_job(service_id):
         raise InvalidRequest("Create job is not allowed: service is inactive ", 403)
 
     data = request.get_json()
-
     data.update({"service": service_id})
+
     try:
         data.update(**get_job_metadata_from_s3(service_id, data["id"]))
     except KeyError:
         raise InvalidRequest({"id": ["Missing data for required field."]}, status_code=400)
 
+    if data.get("valid") != "True":
+        raise InvalidRequest("File is not valid, can't create job", 400)
+
     data["template"] = data.pop("template_id")
+
     template = dao_get_template_by_id(data["template"])
+    template_errors = unarchived_template_schema.validate({"archived": template.archived})
+
+    if template_errors:
+        raise InvalidRequest(template_errors, status_code=400)
 
     job = get_job_from_s3(service_id, data["id"])
     recipient_csv = RecipientCSV(
@@ -170,22 +181,16 @@ def create_job(service_id):
 
         if not is_test_notification:
             check_sms_daily_limit(service, len(recipient_csv))
+            increment_sms_daily_count_send_warnings_if_needed(service, len(recipient_csv))
 
-    if template.template_type == EMAIL_TYPE:
+    elif template.template_type == EMAIL_TYPE:
         check_email_daily_limit(service, len(list(recipient_csv.get_rows())))
+        scheduled_for = datetime.fromisoformat(form.get("scheduled_for")) if form.get("scheduled_for") else None  # noqa: F821
 
-    if data.get("valid") != "True":
-        raise InvalidRequest("File is not valid, can't create job", 400)
-
-    errors = unarchived_template_schema.validate({"archived": template.archived})
-
-    if errors:
-        raise InvalidRequest(errors, status_code=400)
-
-    if template.template_type == SMS_TYPE and not is_test_notification:
-        increment_sms_daily_count_send_warnings_if_needed(service, len(recipient_csv))
-    elif template.template_type == EMAIL_TYPE:
-        increment_email_daily_count_send_warnings_if_needed(service, len(list(recipient_csv.get_rows())))
+        if scheduled_for is None or not scheduled_for.date() > datetime.today().date():
+            increment_email_daily_count_send_warnings_if_needed(
+                authenticated_service, len(list(recipient_csv.get_rows()))  # noqa: F821
+            )
 
     data.update({"template_version": template.version})
 

app/notifications/rest.py

@@ -112,7 +112,7 @@ def send_notification(notification_type: NotificationType):
     )
 
     simulated = simulated_recipient(notification_form["to"], notification_type)
-    if not simulated != api_user.key_type == KEY_TYPE_TEST:
+    if not simulated != api_user.key_type == KEY_TYPE_TEST and notification_type == EMAIL_TYPE:
         check_email_daily_limit(authenticated_service, 1)
 
     check_template_is_for_notification_type(notification_type, template.template_type)

app/v2/notifications/post_notifications.py

@@ -223,7 +223,10 @@ def post_bulk():
 
     if template.template_type == EMAIL_TYPE and api_user.key_type != KEY_TYPE_TEST:
         check_email_daily_limit(authenticated_service, len(list(recipient_csv.get_rows())))
-        increment_email_daily_count_send_warnings_if_needed(authenticated_service, len(list(recipient_csv.get_rows())))
+        scheduled_for = datetime.fromisoformat(form.get("scheduled_for")) if form.get("scheduled_for") else None
+
+        if scheduled_for is None or not scheduled_for.date() > datetime.today().date():
+            increment_email_daily_count_send_warnings_if_needed(authenticated_service, len(list(recipient_csv.get_rows())))
 
     if template.template_type == SMS_TYPE:
         # calculate the number of simulated recipients

tests/app/dao/test_services_dao.py

@@ -46,6 +46,7 @@
 from app.models import (
     EMAIL_TYPE,
     INTERNATIONAL_SMS_TYPE,
+    JOB_STATUS_SCHEDULED,
     KEY_TYPE_NORMAL,
     KEY_TYPE_TEAM,
     KEY_TYPE_TEST,
@@ -68,6 +69,7 @@
     user_folder_permissions,
 )
 from app.schemas import service_schema
+from tests.app.conftest import create_sample_job
 from tests.app.db import (
     create_annual_billing,
     create_api_key,
@@ -1396,7 +1398,7 @@ def create_email_sms_letter_template():
 
 
 class TestServiceEmailLimits:
-    def test_get_email_count_for_service(self, notify_db_session):
+    def test_get_email_count_for_service(self):
         active_user_1 = create_user(email="[email protected]", state="active")
         service = Service(
             name="service_name",
@@ -1436,3 +1438,94 @@ def test_dao_fetch_todays_total_message_count_returns_0_with_yesterday_messages(
             )
         )
         assert fetch_todays_total_message_count(notification.service.id) == 0
+
+    def test_dao_fetch_todays_total_message_count_counts_notifications_in_jobs_scheduled_for_today(
+        self, notify_db, notify_db_session
+    ):
+        service = create_service(service_name="tester")
+        template = create_template(service=service, template_type="email")
+        today = datetime.utcnow().date()
+
+        create_sample_job(
+            notify_db,
+            notify_db_session,
+            service=service,
+            template=template,
+            scheduled_for=today,
+            job_status=JOB_STATUS_SCHEDULED,
+            notification_count=10,
+        )
+        save_notification(
+            create_notification(
+                created_at=today,
+                template=template,
+            )
+        )
+        assert fetch_todays_total_message_count(service.id) == 11
+
+    def test_dao_fetch_todays_total_message_count_counts_notifications_in_jobs_scheduled_for_today_but_not_after_today(
+        self, notify_db, notify_db_session
+    ):
+        service = create_service(service_name="tester")
+        template = create_template(service=service, template_type="email")
+        today = datetime.utcnow().date()
+
+        create_sample_job(
+            notify_db,
+            notify_db_session,
+            service=service,
+            template=template,
+            scheduled_for=today,
+            job_status=JOB_STATUS_SCHEDULED,
+            notification_count=10,
+        )
+        save_notification(
+            create_notification(
+                created_at=today,
+                template=template,
+            )
+        )
+        create_sample_job(
+            notify_db,
+            notify_db_session,
+            service=service,
+            template=template,
+            scheduled_for=today + timedelta(days=1),
+            job_status=JOB_STATUS_SCHEDULED,
+            notification_count=10,
+        )
+
+        assert fetch_todays_total_message_count(service.id) == 11
+
+    def test_dao_fetch_todays_total_message_count_counts_notifications_in_jobs_scheduled_for_today_but_not_before_today(
+        self, notify_db, notify_db_session
+    ):
+        service = create_service(service_name="tester")
+        template = create_template(service=service, template_type="email")
+        today = datetime.utcnow().date()
+
+        create_sample_job(
+            notify_db,
+            notify_db_session,
+            service=service,
+            template=template,
+            scheduled_for=today,
+            job_status=JOB_STATUS_SCHEDULED,
+            notification_count=10,
+        )
+        create_sample_job(
+            notify_db,
+            notify_db_session,
+            service=service,
+            template=template,
+            scheduled_for=today - timedelta(days=1),
+            job_status=JOB_STATUS_SCHEDULED,
+            notification_count=10,
+        )
+        save_notification(
+            create_notification(
+                created_at=today,
+                template=template,
+            )
+        )
+        assert fetch_todays_total_message_count(service.id) == 11

tests/app/job/test_rest.py

@@ -451,6 +451,7 @@ def test_create_job_returns_404_if_template_does_not_exist(client, sample_servic
         "app.job.rest.get_job_metadata_from_s3",
         return_value={
             "template_id": str(sample_service.id),
+            "valid": "True",
         },
     )
     data = {

tests/app/notifications/rest/test_send_notification.py

@@ -466,7 +466,7 @@ def test_should_allow_api_call_if_under_day_limit_regardless_of_type(
             sms_template = create_sample_template(notify_db, notify_db_session, service=service)
             create_sample_notification(notify_db, notify_db_session, template=email_template, service=service)
 
-            data = {"to": sample_user.mobile_number, "template": str(sms_template.id)}
+            data = {"to": sample_user.mobile_number, "template": str(sms_template.id), "valid": "True"}
 
             auth_header = create_authorization_header(service_id=service.id)