From b1ac580e0ab8d8e370cb6265362470b4fa3c2288 Mon Sep 17 00:00:00 2001 From: Ken Tsang Date: Mon, 19 Mar 2018 13:52:01 +0000 Subject: [PATCH] Add celery tasks to handle virus scan passing and failing --- app/celery/letters_pdf_tasks.py | 57 ++++++++++++++++++++-- tests/app/celery/test_letters_pdf_tasks.py | 33 ++++++++++++- 2 files changed, 83 insertions(+), 7 deletions(-) diff --git a/app/celery/letters_pdf_tasks.py b/app/celery/letters_pdf_tasks.py index 4afd13065..b1c329f04 100644 --- a/app/celery/letters_pdf_tasks.py +++ b/app/celery/letters_pdf_tasks.py @@ -1,12 +1,14 @@ import math +from datetime import datetime +from botocore.exceptions import ClientError as BotoClientError from flask import current_app -from notifications_utils.statsd_decorators import statsd from requests import ( post as requests_post, RequestException ) -from botocore.exceptions import ClientError as BotoClientError + +from notifications_utils.statsd_decorators import statsd from app import notify_celery from app.aws import s3 @@ -16,9 +18,15 @@ from app.dao.notifications_dao import ( update_notification_status_by_id, dao_update_notification, dao_get_notifications_by_references, + dao_update_notifications_by_reference, ) -from app.letters.utils import upload_letter_pdf -from app.models import NOTIFICATION_CREATED +from app.letters.utils import ( + delete_pdf_from_letters_scan_bucket, + get_reference_from_filename, + move_scanned_pdf_to_letters_pdf_bucket, + upload_letter_pdf +) +from app.models import NOTIFICATION_CREATED, NOTIFICATION_PERMANENT_FAILURE @notify_celery.task(bind=True, name="create-letters-pdf", max_retries=15, default_retry_delay=300) @@ -133,7 +141,7 @@ def group_letters(letter_pdfs): def letter_in_created_state(filename): # filename looks like '2018-01-13/NOTIFY.ABCDEF1234567890.D.2.C.C.20180113120000.PDF' subfolder = filename.split('/')[0] - ref = filename.split('.')[1] + ref = get_reference_from_filename(filename) notifications = dao_get_notifications_by_references([ref]) if notifications: if notifications[0].status == NOTIFICATION_CREATED: @@ -144,3 +152,42 @@ def letter_in_created_state(filename): notifications[0].status )) return False + + +@notify_celery.task(name='process-letter-scan-passed') +def process_letter_scan_passed(filename): + current_app.logger.info('Virus scan passed: {}'.format(filename)) + move_scanned_pdf_to_letters_pdf_bucket(filename) + reference = get_reference_from_filename(filename) + updated_count = update_letter_pdf_status(reference, NOTIFICATION_CREATED) + + if updated_count != 1: + raise Exception( + "There should only be one letter notification for each reference. Found {} notifications".format( + updated_count + ) + ) + + +@notify_celery.task(name='process-letter-scan-failed') +def process_letter_scan_failed(filename): + current_app.logger.info('Virus scan failed: {}'.format(filename)) + delete_pdf_from_letters_scan_bucket(filename) + reference = get_reference_from_filename(filename) + updated_count = update_letter_pdf_status(reference, NOTIFICATION_PERMANENT_FAILURE) + + if updated_count != 1: + raise Exception( + "There should only be one letter notification for each reference. Found {} notifications".format( + updated_count + ) + ) + + +def update_letter_pdf_status(reference, status): + return dao_update_notifications_by_reference( + references=[reference], + update_dict={ + 'status': status, + 'updated_at': datetime.utcnow() + }) diff --git a/tests/app/celery/test_letters_pdf_tasks.py b/tests/app/celery/test_letters_pdf_tasks.py index 757be3def..903d0ab73 100644 --- a/tests/app/celery/test_letters_pdf_tasks.py +++ b/tests/app/celery/test_letters_pdf_tasks.py @@ -16,10 +16,17 @@ from app.celery.letters_pdf_tasks import ( get_letters_pdf, collate_letter_pdfs_for_day, group_letters, - letter_in_created_state + letter_in_created_state, + process_letter_scan_passed, + process_letter_scan_failed, ) from app.letters.utils import get_letter_pdf_filename -from app.models import Notification, NOTIFICATION_SENDING +from app.models import ( + Notification, + NOTIFICATION_CREATED, + NOTIFICATION_PERMANENT_FAILURE, + NOTIFICATION_SENDING +) from tests.conftest import set_config_values @@ -308,3 +315,25 @@ def test_letter_in_created_state_fails_if_notification_doesnt_exist(sample_notif sample_notification.reference = 'QWERTY1234567890' filename = '2018-01-13/NOTIFY.ABCDEF1234567890.D.2.C.C.20180113120000.PDF' assert letter_in_created_state(filename) is False + + +def test_process_letter_task_check_virus_scan_passed(sample_letter_notification, mocker): + filename = 'NOTIFY.{}'.format(sample_letter_notification.reference) + sample_letter_notification.status = 'pending-virus-check' + mock_move_pdf = mocker.patch('app.celery.letters_pdf_tasks.move_scanned_pdf_to_letters_pdf_bucket') + + process_letter_scan_passed(filename) + + mock_move_pdf.assert_called_once_with(filename) + assert sample_letter_notification.status == NOTIFICATION_CREATED + + +def test_process_letter_task_check_virus_scan_failed(sample_letter_notification, mocker): + filename = 'NOTIFY.{}'.format(sample_letter_notification.reference) + sample_letter_notification.status = 'pending-virus-check' + mock_delete_pdf = mocker.patch('app.celery.letters_pdf_tasks.delete_pdf_from_letters_scan_bucket') + + process_letter_scan_failed(filename) + + mock_delete_pdf.assert_called_once_with(filename) + assert sample_letter_notification.status == NOTIFICATION_PERMANENT_FAILURE