Merge pull request #1716 from alphagov/count-precompiled-letter-pages

Count the number of pages in pre-compiled letters

app/v2/notifications/post_notifications.py

@@ -1,13 +1,16 @@
 import base64
 import functools
+import io
+import math
 
 from flask import request, jsonify, current_app, abort
 
+from notifications_utils.pdf import pdf_page_count, PdfReadError
 from notifications_utils.recipients import try_validate_and_format_phone_number
 
 from app import api_user, authenticated_service
 from app.config import QueueNames
-from app.dao.notifications_dao import update_notification_status_by_reference
+from app.dao.notifications_dao import dao_update_notification, update_notification_status_by_reference
 from app.dao.templates_dao import dao_create_template
 from app.dao.users_dao import get_user_by_id
 from app.letters.utils import upload_letter_pdf
@@ -227,8 +230,12 @@ def process_letter_notification(*, letter_data, api_key, template, reply_to_text
     if precompiled:
         try:
             letter_content = base64.b64decode(letter_data['content'])
+            pages = pdf_page_count(io.BytesIO(letter_content))
         except ValueError:
             raise BadRequestError(message='Cannot decode letter content (invalid base64 encoding)', status_code=400)
+        except PdfReadError:
+            current_app.logger.exception(msg='Invalid PDF received')
+            raise BadRequestError(message='Letter content is not a valid PDF', status_code=400)
 
     should_send = not (api_key.service.research_mode or api_key.key_type == KEY_TYPE_TEST)
 
@@ -243,6 +250,9 @@ def process_letter_notification(*, letter_data, api_key, template, reply_to_text
     if should_send:
         if precompiled:
             upload_letter_pdf(notification, letter_content)
+            pages_per_sheet = 2
+            notification.billable_units = math.ceil(pages / pages_per_sheet)
+            dao_update_notification(notification)
         else:
             create_letters_pdf.apply_async(
                 [str(notification.id)],

tests/app/v2/notifications/test_post_notifications.py

@@ -738,6 +738,7 @@ def test_post_precompiled_letter_with_invalid_base64(client, notify_user, mocker
 def test_post_precompiled_letter_notification_returns_201(client, notify_user, mocker):
     sample_service = create_service(service_permissions=['letter', 'precompiled_letter'])
     s3mock = mocker.patch('app.v2.notifications.post_notifications.upload_letter_pdf')
+    mocker.patch('app.v2.notifications.post_notifications.pdf_page_count', return_value=5)
     data = {
         "reference": "letter-reference",
         "content": "bGV0dGVyLWNvbnRlbnQ="
@@ -754,6 +755,8 @@ def test_post_precompiled_letter_notification_returns_201(client, notify_user, m
 
     notification = Notification.query.first()
 
+    assert notification.billable_units == 3
+
     resp_json = json.loads(response.get_data(as_text=True))
     assert resp_json == {
         'content': {'body': None, 'subject': 'Pre-compiled PDF'},
@@ -767,3 +770,26 @@ def test_post_precompiled_letter_notification_returns_201(client, notify_user, m
         },
         'uri': ANY
     }
+
+
+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