We were getting the page count for the letter before virus scan happened.

This PR moves setting the billale_units for the letter after virus scan has passed.

tests/app/celery/test_letters_pdf_tasks.py

@@ -1,6 +1,7 @@
 from unittest.mock import Mock, call, ANY
 
 import boto3
+from PyPDF2.utils import PdfReadError
 from moto import mock_s3
 from flask import current_app
 from freezegun import freeze_time
@@ -22,7 +23,8 @@ from app.celery.letters_pdf_tasks import (
     process_virus_scan_failed,
     process_virus_scan_error,
     replay_letters_in_error,
-    _sanitise_precomiled_pdf
+    _sanitise_precomiled_pdf,
+    _get_page_count
 )
 from app.letters.utils import get_letter_pdf_filename, ScanErrorType
 from app.models import (
@@ -31,9 +33,12 @@ from app.models import (
     Notification,
     NOTIFICATION_CREATED,
     NOTIFICATION_DELIVERED,
-    NOTIFICATION_VIRUS_SCAN_FAILED,
+    NOTIFICATION_PERMANENT_FAILURE,
     NOTIFICATION_SENDING,
-    NOTIFICATION_TECHNICAL_FAILURE)
+    NOTIFICATION_TECHNICAL_FAILURE,
+    NOTIFICATION_VIRUS_SCAN_FAILED,
+)
+from tests.app.db import create_notification
 
 from tests.conftest import set_config_values
 
@@ -333,14 +338,17 @@ def test_letter_in_created_state_fails_if_notification_doesnt_exist(sample_notif
 
 @freeze_time('2018-01-01 18:00')
 @mock_s3
-@pytest.mark.parametrize('key_type,is_test_letter,noti_status,bucket_config_name,destination_folder', [
-    (KEY_TYPE_NORMAL, False, NOTIFICATION_CREATED, 'LETTERS_PDF_BUCKET_NAME', '2018-01-02/'),
-    (KEY_TYPE_TEST, True, NOTIFICATION_DELIVERED, 'TEST_LETTERS_BUCKET_NAME', '')
+@pytest.mark.parametrize('key_type, noti_status,bucket_config_name,destination_folder', [
+    (KEY_TYPE_NORMAL, NOTIFICATION_CREATED, 'LETTERS_PDF_BUCKET_NAME', '2018-01-02/'),
+    (KEY_TYPE_TEST, NOTIFICATION_DELIVERED, 'TEST_LETTERS_BUCKET_NAME', '')
 ])
 def test_process_letter_task_check_virus_scan_passed(
-    sample_letter_notification, mocker, key_type, is_test_letter, noti_status, bucket_config_name, destination_folder
+    sample_letter_template, mocker, key_type, noti_status, bucket_config_name, destination_folder
 ):
-    filename = 'NOTIFY.{}'.format(sample_letter_notification.reference)
+    letter_notification = create_notification(template=sample_letter_template, billable_units=0,
+                                              status='pending-virus-check', key_type=key_type,
+                                              reference='{} letter'.format(key_type))
+    filename = 'NOTIFY.{}'.format(letter_notification.reference)
     source_bucket_name = current_app.config['LETTERS_SCAN_BUCKET_NAME']
     target_bucket_name = current_app.config[bucket_config_name]
 
@@ -351,14 +359,14 @@ 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'pdf_content')
 
-    sample_letter_notification.status = 'pending-virus-check'
-    sample_letter_notification.key_type = key_type
+    mocker.patch('app.celery.letters_pdf_tasks.pdf_page_count', return_value=1)
     mock_s3upload = mocker.patch('app.celery.letters_pdf_tasks.s3upload')
     mock_sanitise = mocker.patch('app.celery.letters_pdf_tasks._sanitise_precomiled_pdf')
 
     process_virus_scan_passed(filename)
 
-    assert sample_letter_notification.status == noti_status
+    assert letter_notification.status == noti_status
+    assert letter_notification.billable_units == 1
     mock_s3upload.assert_called_once_with(
         filedata=b'pdf_content',
         region='eu-west-1',
@@ -367,11 +375,20 @@ def test_process_letter_task_check_virus_scan_passed(
     )
     mock_sanitise.assert_called_once_with(
         ANY,
-        sample_letter_notification,
+        letter_notification,
         b'pdf_content'
     )
 
 
+def test_get_page_count_set_notification_to_permanent_failure_when_not_pdf(
+        sample_letter_notification
+):
+    with pytest.raises(expected_exception=PdfReadError):
+        _get_page_count(sample_letter_notification, b'pdf_content')
+    updated_notification = Notification.query.filter_by(id=sample_letter_notification.id).first()
+    assert updated_notification.status == NOTIFICATION_PERMANENT_FAILURE
+
+
 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'

tests/app/v2/notifications/test_post_letter_notifications.py

@@ -382,7 +382,6 @@ def test_post_letter_notification_is_delivered_and_has_pdf_uploaded_to_test_lett
 ):
     sample_letter_service = create_service(service_permissions=['letter', 'precompiled_letter'])
     s3mock = mocker.patch('app.v2.notifications.post_notifications.upload_letter_pdf', return_value='test.pdf')
-    mocker.patch('app.v2.notifications.post_notifications.pdf_page_count', return_value=1)
     mock_celery = mocker.patch("app.letters.rest.notify_celery.send_task")
     data = {
         "reference": "letter-reference",
@@ -468,7 +467,6 @@ def test_post_precompiled_letter_notification_returns_201(client, notify_user, m
     sample_service = create_service(service_permissions=['letter', 'precompiled_letter'])
     sample_service.postage = postage
     s3mock = mocker.patch('app.v2.notifications.post_notifications.upload_letter_pdf')
-    mocker.patch('app.v2.notifications.post_notifications.pdf_page_count', return_value=5)
     mocker.patch("app.letters.rest.notify_celery.send_task")
     data = {
         "reference": "letter-reference",
@@ -486,7 +484,7 @@ def test_post_precompiled_letter_notification_returns_201(client, notify_user, m
 
     notification = Notification.query.one()
 
-    assert notification.billable_units == 3
+    assert notification.billable_units == 0
     assert notification.status == NOTIFICATION_PENDING_VIRUS_CHECK
     assert notification.postage == postage
 
@@ -495,26 +493,3 @@ def test_post_precompiled_letter_notification_returns_201(client, notify_user, m
 
     resp_json = json.loads(response.get_data(as_text=True))
     assert resp_json == {'id': str(notification.id), 'reference': 'letter-reference'}
-
-
-def test_post_precompiled_letter_notification_returns_400_with_invalid_pdf(client, notify_user, mocker):
-    sample_service = create_service(service_permissions=['letter', 'precompiled_letter'])
-    s3mock = mocker.patch('app.v2.notifications.post_notifications.upload_letter_pdf')
-    data = {
-        "reference": "letter-reference",
-        "content": "bGV0dGVyLWNvbnRlbnQ="
-    }
-    auth_header = create_authorization_header(service_id=sample_service.id)
-    response = client.post(
-        path="v2/notifications/letter",
-        data=json.dumps(data),
-        headers=[('Content-Type', 'application/json'), auth_header])
-
-    resp_json = json.loads(response.get_data(as_text=True))
-
-    assert response.status_code == 400, response.get_data(as_text=True)
-    assert resp_json['errors'][0]['message'] == 'Letter content is not a valid PDF'
-
-    assert s3mock.called is False
-
-    assert Notification.query.count() == 0