diff --git a/app/celery/letters_pdf_tasks.py b/app/celery/letters_pdf_tasks.py index 7d659c876..3ea1c1bd8 100644 --- a/app/celery/letters_pdf_tasks.py +++ b/app/celery/letters_pdf_tasks.py @@ -1,4 +1,3 @@ -import io import math import base64 from datetime import datetime @@ -9,7 +8,6 @@ from base64 import urlsafe_b64encode from PyPDF2.utils import PdfReadError from botocore.exceptions import ClientError as BotoClientError from flask import current_app -from notifications_utils.pdf import pdf_page_count from requests import ( post as requests_post, RequestException @@ -39,7 +37,8 @@ from app.letters.utils import ( move_failed_pdf, move_scan_to_invalid_pdf_bucket, move_error_pdf_to_scan_bucket, - get_file_names_from_error_bucket + get_file_names_from_error_bucket, + get_page_count, ) from app.models import ( KEY_TYPE_TEST, @@ -211,8 +210,9 @@ def process_virus_scan_passed(self, filename): old_pdf = scan_pdf_object.get()['Body'].read() try: - billable_units = _get_page_count(notification, old_pdf) + billable_units = get_page_count(old_pdf) except PdfReadError: + current_app.logger.exception(msg='Invalid PDF received for notification_id: {}'.format(notification.id)) _move_invalid_letter_and_update_status(notification, filename, scan_pdf_object) return @@ -267,17 +267,6 @@ def process_virus_scan_passed(self, filename): update_notification_status_by_id(notification.id, NOTIFICATION_TECHNICAL_FAILURE) -def _get_page_count(notification, old_pdf): - try: - pages = pdf_page_count(io.BytesIO(old_pdf)) - pages_per_sheet = 2 - billable_units = math.ceil(pages / pages_per_sheet) - return billable_units - except PdfReadError as e: - current_app.logger.exception(msg='Invalid PDF received for notification_id: {}'.format(notification.id)) - raise e - - def _move_invalid_letter_and_update_status(notification, filename, scan_pdf_object): try: move_scan_to_invalid_pdf_bucket(filename) diff --git a/app/letters/utils.py b/app/letters/utils.py index e23542081..2eb7b3131 100644 --- a/app/letters/utils.py +++ b/app/letters/utils.py @@ -1,3 +1,5 @@ +import io +import math from datetime import datetime, timedelta from enum import Enum @@ -5,6 +7,7 @@ import boto3 from flask import current_app from notifications_utils.letter_timings import LETTER_PROCESSING_DEADLINE +from notifications_utils.pdf import pdf_page_count from notifications_utils.s3 import s3upload from notifications_utils.timezones import convert_utc_to_bst @@ -201,3 +204,10 @@ def letter_print_day(created_at): else: print_date = bst_print_datetime.strftime('%d %B').lstrip('0') return 'on {}'.format(print_date) + + +def get_page_count(pdf): + pages = pdf_page_count(io.BytesIO(pdf)) + pages_per_sheet = 2 + billable_units = math.ceil(pages / pages_per_sheet) + return billable_units diff --git a/tests/app/celery/test_letters_pdf_tasks.py b/tests/app/celery/test_letters_pdf_tasks.py index 221162719..5bd72c7f0 100644 --- a/tests/app/celery/test_letters_pdf_tasks.py +++ b/tests/app/celery/test_letters_pdf_tasks.py @@ -414,7 +414,7 @@ def test_process_letter_task_check_virus_scan_passed( s3 = boto3.client('s3', region_name='eu-west-1') s3.put_object(Bucket=source_bucket_name, Key=filename, Body=b'old_pdf') - mock_get_page_count = mocker.patch('app.celery.letters_pdf_tasks._get_page_count', return_value=1) + mock_get_page_count = mocker.patch('app.celery.letters_pdf_tasks.get_page_count', return_value=1) mock_s3upload = mocker.patch('app.celery.letters_pdf_tasks.s3upload') endpoint = 'http://localhost:9999/precompiled/sanitise' with requests_mock.mock() as rmock: @@ -444,10 +444,7 @@ def test_process_letter_task_check_virus_scan_passed( file_location=destination_folder + filename, region='eu-west-1', ) - mock_get_page_count.assert_called_once_with( - letter_notification, - b'old_pdf' - ) + mock_get_page_count.assert_called_once_with(b'old_pdf') @freeze_time('2018-01-01 18:00') @@ -471,7 +468,7 @@ def test_process_letter_task_check_virus_scan_passed_when_sanitise_fails( sample_letter_notification.key_type = key_type mock_move_s3 = mocker.patch('app.letters.utils._move_s3_object') mock_sanitise = mocker.patch('app.celery.letters_pdf_tasks._sanitise_precompiled_pdf', return_value=None) - mock_get_page_count = mocker.patch('app.celery.letters_pdf_tasks._get_page_count', return_value=2) + mock_get_page_count = mocker.patch('app.celery.letters_pdf_tasks.get_page_count', return_value=2) process_virus_scan_passed(filename) @@ -487,9 +484,7 @@ def test_process_letter_task_check_virus_scan_passed_when_sanitise_fails( target_bucket_name, filename ) - mock_get_page_count.assert_called_once_with( - sample_letter_notification, b'pdf_content' - ) + mock_get_page_count.assert_called_once_with(b'pdf_content') @freeze_time('2018-01-01 18:00') @@ -515,7 +510,7 @@ def test_process_letter_task_check_virus_scan_passed_when_redaction_fails( sample_letter_notification.status = NOTIFICATION_PENDING_VIRUS_CHECK sample_letter_notification.key_type = key_type mock_copy_s3 = mocker.patch('app.letters.utils._copy_s3_object') - mocker.patch('app.celery.letters_pdf_tasks._get_page_count', return_value=2) + mocker.patch('app.celery.letters_pdf_tasks.get_page_count', return_value=2) endpoint = 'http://localhost:9999/precompiled/sanitise' with requests_mock.mock() as rmock: @@ -564,15 +559,13 @@ def test_process_letter_task_check_virus_scan_passed_when_file_cannot_be_opened( sample_letter_notification.key_type = key_type mock_move_s3 = mocker.patch('app.letters.utils._move_s3_object') - mock_get_page_count = mocker.patch('app.celery.letters_pdf_tasks._get_page_count', side_effect=PdfReadError) + mock_get_page_count = mocker.patch('app.celery.letters_pdf_tasks.get_page_count', side_effect=PdfReadError) mock_sanitise = mocker.patch('app.celery.letters_pdf_tasks._sanitise_precompiled_pdf') process_virus_scan_passed(filename) mock_sanitise.assert_not_called() - mock_get_page_count.assert_called_once_with( - sample_letter_notification, b'pdf_content' - ) + mock_get_page_count.assert_called_once_with(b'pdf_content') mock_move_s3.assert_called_once_with( source_bucket_name, filename, target_bucket_name, filename @@ -598,7 +591,7 @@ def test_process_virus_scan_passed_logs_error_and_sets_tech_failure_if_s3_error_ s3 = boto3.client('s3', region_name='eu-west-1') s3.put_object(Bucket=source_bucket_name, Key=filename, Body=b'pdf_content') - mocker.patch('app.celery.letters_pdf_tasks._get_page_count', return_value=1) + mocker.patch('app.celery.letters_pdf_tasks.get_page_count', return_value=1) error_response = { 'Error': {