From 081543a2a93d76e015ea3686683be30715d63899 Mon Sep 17 00:00:00 2001 From: Katie Smith Date: Thu, 5 Sep 2019 16:45:48 +0100 Subject: [PATCH] Refactor out function to get page count This has been moved to the letters utils file since it will be used in more than one place. The notification parameter has been removed so that the function can be used when we don't have a notification id. --- app/celery/letters_pdf_tasks.py | 19 ++++-------------- app/letters/utils.py | 10 ++++++++++ tests/app/celery/test_letters_pdf_tasks.py | 23 ++++++++-------------- 3 files changed, 22 insertions(+), 30 deletions(-) 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': {