Added a test

By added the notification_type to the filter we get a better performance on the query. Especially when selecting for NotificationHistory.

app/celery/letters_pdf_tasks.py

@@ -42,6 +42,7 @@ from app.models import (
     NOTIFICATION_TECHNICAL_FAILURE,
     NOTIFICATION_VALIDATION_FAILED,
     NOTIFICATION_VIRUS_SCAN_FAILED,
+    LETTER_TYPE
 )
 from app.cronitor import cronitor
 
@@ -216,7 +217,7 @@ def group_letters(letter_pdfs):
 def sanitise_letter(self, filename):
     try:
         reference = get_reference_from_filename(filename)
-        notification = dao_get_notification_by_reference(reference)
+        notification = dao_get_notification_by_reference(reference=reference, notification_type=LETTER_TYPE)
 
         current_app.logger.info('Notification ID {} Virus scan passed: {}'.format(notification.id, filename))
 
@@ -352,7 +353,7 @@ def _move_invalid_letter_and_update_status(
 def process_virus_scan_failed(filename):
     move_failed_pdf(filename, ScanErrorType.FAILURE)
     reference = get_reference_from_filename(filename)
-    notification = dao_get_notification_by_reference(reference)
+    notification = dao_get_notification_by_reference(reference=reference, notification_type=LETTER_TYPE)
     updated_count = update_letter_pdf_status(reference, NOTIFICATION_VIRUS_SCAN_FAILED, billable_units=0)
 
     if updated_count != 1:
@@ -371,7 +372,7 @@ def process_virus_scan_failed(filename):
 def process_virus_scan_error(filename):
     move_failed_pdf(filename, ScanErrorType.ERROR)
     reference = get_reference_from_filename(filename)
-    notification = dao_get_notification_by_reference(reference)
+    notification = dao_get_notification_by_reference(reference=reference, notification_type=LETTER_TYPE)
     updated_count = update_letter_pdf_status(reference, NOTIFICATION_TECHNICAL_FAILURE, billable_units=0)
 
     if updated_count != 1:

app/celery/process_ses_receipts_tasks.py

@@ -10,7 +10,7 @@ from app import notify_celery, statsd_client
 from app.config import QueueNames
 from app.clients.email.aws_ses import get_aws_responses
 from app.dao import notifications_dao
-from app.models import NOTIFICATION_SENDING, NOTIFICATION_PENDING
+from app.models import NOTIFICATION_SENDING, NOTIFICATION_PENDING, EMAIL_TYPE
 
 from app.notifications.notifications_ses_callback import (
     determine_notification_bounce_type,
@@ -39,7 +39,9 @@ def process_ses_results(self, response):
         reference = ses_message['mail']['messageId']
 
         try:
-            notification = notifications_dao.dao_get_notification_or_history_by_reference(reference=reference)
+            notification = notifications_dao.dao_get_notification_or_history_by_reference(
+                reference=reference, notification_type=EMAIL_TYPE
+            )
         except NoResultFound:
             message_time = iso8601.parse_date(ses_message['mail']['timestamp']).replace(tzinfo=None)
             if datetime.utcnow() - message_time < timedelta(minutes=5):

app/celery/tasks.py

@@ -536,7 +536,7 @@ def update_letter_notification(filename, temporary_failures, update):
 
 
 def check_billable_units(notification_update):
-    notification = dao_get_notification_or_history_by_reference(notification_update.reference)
+    notification = dao_get_notification_or_history_by_reference(notification_update.reference, LETTER_TYPE)
 
     if int(notification_update.page_count) != notification.billable_units:
         msg = 'Notification with id {} has {} billable_units but DVLA says page count is {}'.format(

app/dao/notifications_dao.py

@@ -650,33 +650,29 @@ def dao_get_notifications_by_recipient_or_reference(
 
 
 @statsd(namespace="dao")
-def dao_get_notification_by_reference(reference):
+def dao_get_notification_by_reference(reference, notification_type):
     return Notification.query.filter(
-        Notification.reference == reference
+        Notification.reference == reference,
+        Notification.notification_type == notification_type
     ).one()
 
 
 @statsd(namespace="dao")
-def dao_get_notification_or_history_by_reference(reference):
+def dao_get_notification_or_history_by_reference(reference, notification_type):
     try:
         # This try except is necessary because in test keys and research mode does not create notification history.
         # Otherwise we could just search for the NotificationHistory object
         return Notification.query.filter(
-            Notification.reference == reference
+            Notification.reference == reference,
+            Notification.notification_type == notification_type
         ).one()
     except NoResultFound:
         return NotificationHistory.query.filter(
-            NotificationHistory.reference == reference
+            NotificationHistory.reference == reference,
+            NotificationHistory.notification_type == notification_type
         ).one()
 
 
-@statsd(namespace="dao")
-def dao_get_notifications_by_references(references):
-    return Notification.query.filter(
-        Notification.reference.in_(references)
-    ).all()
-
-
 @statsd(namespace="dao")
 def dao_created_scheduled_notification(scheduled_notification):
     db.session.add(scheduled_notification)

app/notifications/notifications_ses_callback.py

@@ -5,7 +5,7 @@ from app.dao.notifications_dao import dao_get_notification_or_history_by_referen
 from app.dao.service_callback_api_dao import (
     get_service_delivery_status_callback_api_for_service, get_service_complaint_callback_api_for_service
 )
-from app.models import Complaint
+from app.models import Complaint, EMAIL_TYPE
 from app.celery.service_callback_tasks import (
     send_delivery_status_to_service,
     send_complaint_to_service,
@@ -33,7 +33,7 @@ def handle_complaint(ses_message):
     except KeyError as e:
         current_app.logger.exception("Complaint from SES failed to get reference from message", e)
         return
-    notification = dao_get_notification_or_history_by_reference(reference)
+    notification = dao_get_notification_or_history_by_reference(reference, EMAIL_TYPE)
     ses_complaint = ses_message.get('complaint', None)
 
     complaint = Complaint(

tests/app/dao/notification_dao/test_notification_dao.py

@@ -28,7 +28,6 @@ from app.dao.notifications_dao import (
     update_notification_status_by_id,
     update_notification_status_by_reference,
     dao_get_notification_by_reference,
-    dao_get_notifications_by_references,
     dao_get_notification_or_history_by_reference,
     notifications_not_yet_sent,
 )
@@ -1613,7 +1612,7 @@ def test_dao_update_notifications_by_reference_updates_history_when_one_of_two_n
 
 def test_dao_get_notification_by_reference_with_one_match_returns_notification(sample_letter_template, notify_db):
     create_notification(template=sample_letter_template, reference='REF1')
-    notification = dao_get_notification_by_reference('REF1')
+    notification = dao_get_notification_by_reference('REF1', 'letter')
 
     assert notification.reference == 'REF1'
 
@@ -1623,30 +1622,25 @@ def test_dao_get_notification_by_reference_with_multiple_matches_raises_error(sa
     create_notification(template=sample_letter_template, reference='REF1')
 
     with pytest.raises(SQLAlchemyError):
-        dao_get_notification_by_reference('REF1')
+        dao_get_notification_by_reference('REF1', 'letter')
 
 
 def test_dao_get_notification_by_reference_with_no_matches_raises_error(notify_db):
     with pytest.raises(SQLAlchemyError):
-        dao_get_notification_by_reference('REF1')
+        dao_get_notification_by_reference('REF1', 'email')
 
 
-def test_dao_get_notifications_by_references(sample_template):
-    create_notification(template=sample_template, reference='noref')
-    notification_1 = create_notification(template=sample_template, reference='ref')
-    notification_2 = create_notification(template=sample_template, reference='ref')
-
-    notifications = dao_get_notifications_by_references(['ref'])
-    assert len(notifications) == 2
-    assert notifications[0].id in [notification_1.id, notification_2.id]
-    assert notifications[1].id in [notification_1.id, notification_2.id]
+def test_dao_get_notification_by_reference_with_no_matches_for_type_raises_error(sample_email_template):
+    create_notification(template=sample_email_template, reference='REF1')
+    with pytest.raises(SQLAlchemyError):
+        dao_get_notification_by_reference('REF1', 'letter')
 
 
 def test_dao_get_notification_or_history_by_reference_with_one_match_returns_notification(
         sample_letter_template
 ):
     create_notification(template=sample_letter_template, reference='REF1')
-    notification = dao_get_notification_or_history_by_reference('REF1')
+    notification = dao_get_notification_or_history_by_reference('REF1', 'letter')
 
     assert notification.reference == 'REF1'
 
@@ -1658,12 +1652,18 @@ def test_dao_get_notification_or_history_by_reference_with_multiple_matches_rais
     create_notification(template=sample_letter_template, reference='REF1')
 
     with pytest.raises(SQLAlchemyError):
-        dao_get_notification_or_history_by_reference('REF1')
+        dao_get_notification_or_history_by_reference('REF1', 'letter')
 
 
-def test_dao_get_notification_or_history_by_reference_with_no_matches_raises_error(notify_db):
+def test_dao_get_notification_or_history_by_reference_with_no_matches_raises_error(sample_letter_template):
+    create_notification(template=sample_letter_template, reference='REF1')
     with pytest.raises(SQLAlchemyError):
-        dao_get_notification_or_history_by_reference('REF1')
+        dao_get_notification_or_history_by_reference('REF1', 'email')
+
+
+def test_dao_get_notification_or_history_by_reference_with_no_matches_for_type_raises_error(notify_db):
+    with pytest.raises(SQLAlchemyError):
+        dao_get_notification_or_history_by_reference('REF1', 'email')
 
 
 @pytest.mark.parametrize("notification_type",