diff --git a/app/dao/notifications_dao.py b/app/dao/notifications_dao.py index 564ab6acf..a80cdbefe 100644 --- a/app/dao/notifications_dao.py +++ b/app/dao/notifications_dao.py @@ -39,6 +39,7 @@ from app.models import ( NOTIFICATION_DELIVERED, NOTIFICATION_SENDING, NOTIFICATION_PENDING, + NOTIFICATION_PENDING_VIRUS_CHECK, NOTIFICATION_TECHNICAL_FAILURE, NOTIFICATION_TEMPORARY_FAILURE, NOTIFICATION_PERMANENT_FAILURE, @@ -145,13 +146,14 @@ def _update_notification_status(notification, status): @statsd(namespace="dao") @transactional def update_notification_status_by_id(notification_id, status, sent_by=None): - notification = Notification.query.with_lockmode("update").filter( + notification = Notification.query.with_for_update().filter( Notification.id == notification_id, or_( Notification.status == NOTIFICATION_CREATED, Notification.status == NOTIFICATION_SENDING, Notification.status == NOTIFICATION_PENDING, - Notification.status == NOTIFICATION_SENT + Notification.status == NOTIFICATION_SENT, + Notification.status == NOTIFICATION_PENDING_VIRUS_CHECK )).first() if not notification: diff --git a/tests/app/dao/notification_dao/test_notification_dao.py b/tests/app/dao/notification_dao/test_notification_dao.py index ad3a9aa41..50c100192 100644 --- a/tests/app/dao/notification_dao/test_notification_dao.py +++ b/tests/app/dao/notification_dao/test_notification_dao.py @@ -141,6 +141,14 @@ def test_should_update_status_by_id_if_created(notify_db, notify_db_session): assert updated.status == 'failed' +def test_should_update_status_by_id_if_pending_virus_check(notify_db, notify_db_session): + notification = sample_notification(notify_db, notify_db_session, status='pending-virus-check') + assert Notification.query.get(notification.id).status == 'pending-virus-check' + updated = update_notification_status_by_id(notification.id, 'cancelled') + assert Notification.query.get(notification.id).status == 'cancelled' + assert updated.status == 'cancelled' + + def test_should_update_status_by_id_and_set_sent_by(notify_db, notify_db_session): notification = sample_notification(notify_db, notify_db_session, status='sending')