From 63fe21020216afe134272612706b82eccdc80b27 Mon Sep 17 00:00:00 2001
From: Pea Tyczynska <pea.tyczynska@digital.cabinet-office.gov.uk>
Date: Thu, 4 Apr 2019 18:38:31 +0100
Subject: [PATCH] Also show overlay on document download of failed letter

And of course test this new flex
---
 app/letters/utils.py            |  1 -
 app/template/rest.py            | 16 +++++++++++-----
 tests/app/template/test_rest.py | 19 ++++++++++++-------
 3 files changed, 23 insertions(+), 13 deletions(-)

diff --git a/app/letters/utils.py b/app/letters/utils.py
index c7e558244..786edbc57 100644
--- a/app/letters/utils.py
+++ b/app/letters/utils.py
@@ -136,7 +136,6 @@ def get_letter_pdf(notification):
 
     s3 = boto3.resource('s3')
     bucket = s3.Bucket(bucket_name)
-
     item = next(x for x in bucket.objects.filter(Prefix=prefix))
 
     obj = s3.Object(
diff --git a/app/template/rest.py b/app/template/rest.py
index be09cb74e..914520cda 100644
--- a/app/template/rest.py
+++ b/app/template/rest.py
@@ -251,10 +251,9 @@ def preview_letter_template_by_notification_id(service_id, notification_id, file
 
             if request.args.get('overlay'):
                 content = pdf_page
-                url = '{}/precompiled/overlay.png{}{}'.format(
+                url = '{}/precompiled/overlay.png{}'.format(
                     current_app.config['TEMPLATE_PREVIEW_API_HOST'],
                     '?invert=1',
-                    '&hide_notify=true' if page_number == '1' else ''
                 )
             else:
                 url = '{}/precompiled-preview.png{}'.format(
@@ -262,7 +261,14 @@ def preview_letter_template_by_notification_id(service_id, notification_id, file
                     '?hide_notify=true' if page_number == '1' else ''
                 )
 
-            content = _get_png_preview(url, content, notification.id, json=False)
+            content = _get_png_preview_or_overlaid_pdf(url, content, notification.id, json=False)
+        elif file_type == 'pdf' and request.args.get('overlay'):
+            content = pdf_file
+            url = '{}/precompiled/overlay.pdf{}'.format(
+                current_app.config['TEMPLATE_PREVIEW_API_HOST'],
+                '?invert=1',
+            )
+            content = _get_png_preview_or_overlaid_pdf(url, content, notification.id, json=False)
     else:
 
         template_for_letter_print = {
@@ -287,12 +293,12 @@ def preview_letter_template_by_notification_id(service_id, notification_id, file
             file_type,
             '?page={}'.format(page) if page else ''
         )
-        content = _get_png_preview(url, data, notification.id, json=True)
+        content = _get_png_preview_or_overlaid_pdf(url, data, notification.id, json=True)
 
     return jsonify({"content": content})
 
 
-def _get_png_preview(url, data, notification_id, json=True):
+def _get_png_preview_or_overlaid_pdf(url, data, notification_id, json=True):
     if json:
         resp = requests_post(
             url,
diff --git a/tests/app/template/test_rest.py b/tests/app/template/test_rest.py
index 030ac6674..dc7dbbd2c 100644
--- a/tests/app/template/test_rest.py
+++ b/tests/app/template/test_rest.py
@@ -1099,15 +1099,20 @@ def test_preview_letter_template_precompiled_s3_error(
 
 
 @pytest.mark.parametrize(
-    "post_url, overlay",
-    [('precompiled-preview.png', None), ('precompiled/overlay.png?invert=1', 1)]
+    "filetype, post_url, overlay",
+    [
+        ('png', 'precompiled-preview.png', None),
+        ('png', 'precompiled/overlay.png?invert=1', 1),
+        ('pdf', 'precompiled/overlay.pdf?invert=1', 1)
+    ]
 )
-def test_preview_letter_template_precompiled_png_file_type(
+def test_preview_letter_template_precompiled_png_file_type_or_pdf_with_overlay(
         notify_api,
         client,
         admin_request,
         sample_service,
         mocker,
+        filetype,
         post_url,
         overlay
 ):
@@ -1127,7 +1132,7 @@ def test_preview_letter_template_precompiled_png_file_type(
         with requests_mock.Mocker() as request_mock:
 
             pdf_content = b'\x00\x01'
-            png_content = b'\x00\x02'
+            expected_returned_content = b'\x00\x02'
 
             mock_get_letter_pdf = mocker.patch('app.template.rest.get_letter_pdf', return_value=pdf_content)
 
@@ -1135,7 +1140,7 @@ def test_preview_letter_template_precompiled_png_file_type(
 
             mock_post = request_mock.post(
                 'http://localhost/notifications-template-preview/{}'.format(post_url),
-                content=png_content,
+                content=expected_returned_content,
                 headers={'X-pdf-page-count': '1'},
                 status_code=200
             )
@@ -1144,14 +1149,14 @@ def test_preview_letter_template_precompiled_png_file_type(
                 'template.preview_letter_template_by_notification_id',
                 service_id=notification.service_id,
                 notification_id=notification.id,
-                file_type='png',
+                file_type=filetype,
                 overlay=overlay,
             )
 
             with pytest.raises(ValueError):
                 mock_post.last_request.json()
             assert mock_get_letter_pdf.called_once_with(notification)
-            assert base64.b64decode(resp['content']) == png_content
+            assert base64.b64decode(resp['content']) == expected_returned_content
 
 
 @pytest.mark.parametrize('page_number,expect_preview_url', [