diff --git a/app/celery/letters_pdf_tasks.py b/app/celery/letters_pdf_tasks.py index 24c7347ae..d70265791 100644 --- a/app/celery/letters_pdf_tasks.py +++ b/app/celery/letters_pdf_tasks.py @@ -1,4 +1,5 @@ from flask import current_app +import math from requests import ( post as requests_post, RequestException @@ -9,7 +10,11 @@ from botocore.exceptions import ClientError as BotoClientError from app import notify_celery from app.aws import s3 from app.config import QueueNames -from app.dao.notifications_dao import get_notification_by_id, update_notification_status_by_id +from app.dao.notifications_dao import ( + get_notification_by_id, + update_notification_status_by_id, + dao_update_notification +) from app.statsd_decorators import statsd @@ -19,7 +24,7 @@ def create_letters_pdf(self, notification_id): try: notification = get_notification_by_id(notification_id, _raise=True) - pdf_data = get_letters_pdf( + pdf_data, billable_units = get_letters_pdf( notification.template, contact_block=notification.reply_to_text, org_id=notification.service.dvla_organisation.id, @@ -28,6 +33,14 @@ def create_letters_pdf(self, notification_id): current_app.logger.info("PDF Letter {} reference {} created at {}, {} bytes".format( notification.id, notification.reference, notification.created_at, len(pdf_data))) s3.upload_letters_pdf(reference=notification.reference, crown=notification.service.crown, filedata=pdf_data) + + notification.billable_units = billable_units + dao_update_notification(notification) + + current_app.logger.info( + 'Letter notification reference {reference}: billable units set to {billable_units}'.format( + reference=str(notification.reference), billable_units=billable_units)) + except (RequestException, BotoClientError): try: current_app.logger.exception( @@ -62,4 +75,7 @@ def get_letters_pdf(template, contact_block, org_id, values): ) resp.raise_for_status() - return resp.content + pages_per_sheet = 2 + billable_units = math.ceil(int(resp.headers.get("X-pdf-page-count", 0)) / pages_per_sheet) + + return resp.content, billable_units diff --git a/tests/app/celery/test_letters_pdf_tasks.py b/tests/app/celery/test_letters_pdf_tasks.py index 0af5c78ab..a414fa2db 100644 --- a/tests/app/celery/test_letters_pdf_tasks.py +++ b/tests/app/celery/test_letters_pdf_tasks.py @@ -10,6 +10,7 @@ from app.celery.letters_pdf_tasks import ( create_letters_pdf, get_letters_pdf, ) +from app.models import Notification from tests.conftest import set_config_values @@ -46,8 +47,36 @@ def test_get_letters_pdf_calls_notifications_template_preview_service_correctly( } +@pytest.mark.parametrize('page_count,expected_billable_units', [ + ('1', 1), + ('2', 1), + ('3', 2) +]) +def test_get_letters_pdf_calculates_billing_units( + notify_api, mocker, client, sample_letter_template, page_count, expected_billable_units): + contact_block = 'Mr Foo,\n1 Test Street,\nLondon\nN1' + dvla_org_id = '002' + + with set_config_values(notify_api, { + 'TEMPLATE_PREVIEW_API_HOST': 'http://localhost/notifications-template-preview', + 'TEMPLATE_PREVIEW_API_KEY': 'test-key' + }): + with requests_mock.Mocker() as request_mock: + request_mock.post( + 'http://localhost/notifications-template-preview/print.pdf', + content=b'\x00\x01', + headers={'X-pdf-page-count': page_count}, + status_code=200 + ) + + _, billable_units = get_letters_pdf( + sample_letter_template, contact_block=contact_block, org_id=dvla_org_id, values=None) + + assert billable_units == expected_billable_units + + def test_create_letters_pdf_calls_upload_letters_pdf(mocker, sample_letter_notification): - mocker.patch('app.celery.letters_pdf_tasks.get_letters_pdf', return_value=b'\x00\x01') + mocker.patch('app.celery.letters_pdf_tasks.get_letters_pdf', return_value=(b'\x00\x01', '1')) mock_s3 = mocker.patch('app.celery.tasks.s3.upload_letters_pdf') create_letters_pdf(sample_letter_notification.id) @@ -59,6 +88,15 @@ def test_create_letters_pdf_calls_upload_letters_pdf(mocker, sample_letter_notif ) +def test_create_letters_pdf_sets_billable_units(mocker, sample_letter_notification): + mocker.patch('app.celery.letters_pdf_tasks.get_letters_pdf', return_value=(b'\x00\x01', 1)) + mocker.patch('app.celery.tasks.s3.upload_letters_pdf') + + create_letters_pdf(sample_letter_notification.id) + noti = Notification.query.filter(Notification.reference == sample_letter_notification.reference).one() + assert noti.billable_units == 1 + + def test_create_letters_pdf_non_existent_notification(notify_api, mocker, fake_uuid): with pytest.raises(expected_exception=NoResultFound): create_letters_pdf(fake_uuid) @@ -75,7 +113,7 @@ def test_create_letters_pdf_handles_request_errors(mocker, sample_letter_notific def test_create_letters_pdf_handles_s3_errors(mocker, sample_letter_notification): - mocker.patch('app.celery.letters_pdf_tasks.get_letters_pdf') + mocker.patch('app.celery.letters_pdf_tasks.get_letters_pdf', return_value=(b'\x00\x01', 1)) mock_s3 = mocker.patch('app.celery.tasks.s3.upload_letters_pdf', side_effect=ClientError({}, 'operation_name')) mock_retry = mocker.patch('app.celery.letters_pdf_tasks.create_letters_pdf.retry')