From dd9e2856848bfa0d8f236f59e62b072a7fc2e87c Mon Sep 17 00:00:00 2001
From: Pea Tyczynska <pea.tyczynska@digital.cabinet-office.gov.uk>
Date: Wed, 10 Apr 2019 16:08:08 +0100
Subject: [PATCH] Refactor preview_letter_template_by_notification_id

---
 app/template/rest.py            | 45 +++++++++++++++------------------
 tests/app/template/test_rest.py |  2 +-
 2 files changed, 21 insertions(+), 26 deletions(-)

diff --git a/app/template/rest.py b/app/template/rest.py
index 914520cda..2c3755b74 100644
--- a/app/template/rest.py
+++ b/app/template/rest.py
@@ -235,13 +235,23 @@ def preview_letter_template_by_notification_id(service_id, notification_id, file
             )
 
         content = base64.b64encode(pdf_file).decode('utf-8')
+        overlay = request.args.get('overlay')
+        page_number = page if page else "1"
+
+        if overlay:
+            path = '/precompiled/overlay.{}'.format(file_type)
+            query_string = '?invert=1'
+            content = pdf_file
+        elif file_type == 'png':
+            path = '/precompiled-preview.png'
+            query_string = '?hide_notify=true' if page_number == '1' else ''
+        else:
+            path = None
 
         if file_type == 'png':
             try:
-                page_number = page if page else "1"
-
                 pdf_page = extract_page_from_pdf(BytesIO(pdf_file), int(page_number) - 1)
-                content = base64.b64encode(pdf_page).decode('utf-8')
+                content = pdf_page if overlay else base64.b64encode(pdf_page).decode('utf-8')
             except PdfReadError as e:
                 raise InvalidRequest(
                     'Error extracting requested page from PDF file for notification_id {} type {} {}'.format(
@@ -249,26 +259,11 @@ def preview_letter_template_by_notification_id(service_id, notification_id, file
                     status_code=500
                 )
 
-            if request.args.get('overlay'):
-                content = pdf_page
-                url = '{}/precompiled/overlay.png{}'.format(
-                    current_app.config['TEMPLATE_PREVIEW_API_HOST'],
-                    '?invert=1',
-                )
-            else:
-                url = '{}/precompiled-preview.png{}'.format(
-                    current_app.config['TEMPLATE_PREVIEW_API_HOST'],
-                    '?hide_notify=true' if page_number == '1' else ''
-                )
-
-            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)
+        if path:
+            url = current_app.config['TEMPLATE_PREVIEW_API_HOST'] + path + query_string
+            response_content = _get_png_preview_or_overlaid_pdf(url, content, notification.id, json=False)
+        else:
+            response_content = content
     else:
 
         template_for_letter_print = {
@@ -293,9 +288,9 @@ def preview_letter_template_by_notification_id(service_id, notification_id, file
             file_type,
             '?page={}'.format(page) if page else ''
         )
-        content = _get_png_preview_or_overlaid_pdf(url, data, notification.id, json=True)
+        response_content = _get_png_preview_or_overlaid_pdf(url, data, notification.id, json=True)
 
-    return jsonify({"content": content})
+    return jsonify({"content": response_content})
 
 
 def _get_png_preview_or_overlaid_pdf(url, data, notification_id, json=True):
diff --git a/tests/app/template/test_rest.py b/tests/app/template/test_rest.py
index dc7dbbd2c..b0ff45ee8 100644
--- a/tests/app/template/test_rest.py
+++ b/tests/app/template/test_rest.py
@@ -1192,7 +1192,7 @@ def test_preview_letter_template_precompiled_png_file_type_hide_notify_tag_only_
 
         mocker.patch('app.template.rest.get_letter_pdf', return_value=pdf_content)
         mocker.patch('app.template.rest.extract_page_from_pdf', return_value=png_content)
-        mock_get_png_preview = mocker.patch('app.template.rest._get_png_preview', return_value=encoded)
+        mock_get_png_preview = mocker.patch('app.template.rest._get_png_preview_or_overlaid_pdf', return_value=encoded)
 
         admin_request.get(
             'template.preview_letter_template_by_notification_id',