Add a retry mechanism when fetching email attachements

app/delivery/send_to_providers.py

@@ -1,7 +1,6 @@
 import base64
 import os
 import re
-import urllib.request
 from datetime import datetime
 from typing import Dict
 from uuid import UUID
@@ -17,6 +16,8 @@
     SMSMessageTemplate,
 )
 from unidecode import unidecode
+from urllib3 import PoolManager
+from urllib3.util import Retry
 
 from app import bounce_rate_client, clients, document_download_client, statsd_client
 from app.celery.research_mode_tasks import send_email_response, send_sms_response
@@ -231,27 +232,29 @@ def send_email_to_provider(notification: Notification):
             check_file_url(personalisation_data[key]["document"], notification.id)
             sending_method = personalisation_data[key]["document"].get("sending_method")
             direct_file_url = personalisation_data[key]["document"]["direct_file_url"]
+            filename = personalisation_data[key]["document"].get("filename")
+            mime_type = personalisation_data[key]["document"].get("mime_type")
             document_id = personalisation_data[key]["document"]["id"]
             scan_verdict_response = document_download_client.check_scan_verdict(service.id, document_id, sending_method)
             check_for_malware_errors(scan_verdict_response.status_code, notification)
             current_app.logger.info(f"scan_verdict for document_id {document_id} is {scan_verdict_response.json()}")
             if sending_method == "attach":
                 try:
-                    req = urllib.request.Request(direct_file_url)
-                    with urllib.request.urlopen(req) as response:
-                        buffer = response.read()
-                        filename = personalisation_data[key]["document"].get("filename")
-                        mime_type = personalisation_data[key]["document"].get("mime_type")
-                        attachments.append(
-                            {
-                                "name": filename,
-                                "data": buffer,
-                                "mime_type": mime_type,
-                            }
-                        )
+                    retries = Retry(total=5)
+                    http = PoolManager(retries=retries)
+
+                    response = http.request("GET", url=direct_file_url)
+                    attachments.append(
+                        {
+                            "name": filename,
+                            "data": response.data,
+                            "mime_type": mime_type,
+                        }
+                    )
                 except Exception as e:
                     current_app.logger.error(f"Could not download and attach {direct_file_url}\nException: {e}")
-                del personalisation_data[key]
+                    del personalisation_data[key]
+
             else:
                 personalisation_data[key] = personalisation_data[key]["document"]["url"]
 

tests/app/delivery/test_send_to_providers.py

@@ -1,8 +1,9 @@
 import uuid
 from collections import namedtuple
 from datetime import datetime
+from http.client import HTTPMessage
 from unittest import TestCase
-from unittest.mock import ANY, MagicMock, call
+from unittest.mock import ANY, MagicMock, Mock, call
 
 import pytest
 from flask import current_app
@@ -884,6 +885,82 @@ def test_send_email_to_provider_should_format_email_address(sample_email_notific
     )
 
 
+def test_file_attachment_retry(mocker, notify_db, notify_db_session):
+    template = create_sample_email_template(notify_db, notify_db_session, content="Here is your ((file))")
+
+    class mock_response:
+        status_code = 200
+
+        def json():
+            return {"av-status": "clean"}
+
+    mocker.patch("app.delivery.send_to_providers.document_download_client.check_scan_verdict", return_value=mock_response)
+
+    personalisation = {
+        "file": document_download_response(
+            {
+                "direct_file_url": "http://foo.bar/direct_file_url",
+                "url": "http://foo.bar/url",
+                "mime_type": "application/pdf",
+            }
+        )
+    }
+    personalisation["file"]["document"]["sending_method"] = "attach"
+    personalisation["file"]["document"]["filename"] = "file.txt"
+
+    db_notification = save_notification(create_notification(template=template, personalisation=personalisation))
+
+    mocker.patch("app.delivery.send_to_providers.statsd_client")
+    mocker.patch("app.aws_ses_client.send_email", return_value="reference")
+
+    getconn_mock = mocker.patch("urllib3.connectionpool.HTTPConnectionPool._new_conn")
+    getconn_mock.return_value.getresponse.side_effect = [
+        Mock(status=500, msg=HTTPMessage()),
+        Mock(status=429, msg=HTTPMessage()),
+        Mock(status=400, msg=HTTPMessage()),
+        Mock(status=404, msg=HTTPMessage()),
+        Mock(status=200, msg=HTTPMessage()),
+    ]
+
+    mock_logger = mocker.patch("app.delivery.send_to_providers.current_app.logger.error")
+    send_to_providers.send_email_to_provider(db_notification)
+    assert mock_logger.call_count == 0
+
+
+def test_file_attachment_max_retries(mocker, notify_db, notify_db_session):
+    template = create_sample_email_template(notify_db, notify_db_session, content="Here is your ((file))")
+
+    class mock_response:
+        status_code = 200
+
+        def json():
+            return {"av-status": "clean"}
+
+    mocker.patch("app.delivery.send_to_providers.document_download_client.check_scan_verdict", return_value=mock_response)
+
+    personalisation = {
+        "file": document_download_response(
+            {
+                "direct_file_url": "http://foo.bar/direct_file_url",
+                "url": "http://foo.bar/url",
+                "mime_type": "application/pdf",
+            }
+        )
+    }
+    personalisation["file"]["document"]["sending_method"] = "attach"
+    personalisation["file"]["document"]["filename"] = "file.txt"
+
+    db_notification = save_notification(create_notification(template=template, personalisation=personalisation))
+
+    mocker.patch("app.delivery.send_to_providers.statsd_client")
+    mocker.patch("app.aws_ses_client.send_email", return_value="reference")
+
+    mock_logger = mocker.patch("app.delivery.send_to_providers.current_app.logger.error")
+    send_to_providers.send_email_to_provider(db_notification)
+    assert mock_logger.call_count == 1
+    assert "Max retries exceeded" in mock_logger.call_args[0][0]
+
+
 @pytest.mark.parametrize(
     "filename_attribute_present, filename, expected_filename",
     [
@@ -930,28 +1007,26 @@ def json():
 
     statsd_mock = mocker.patch("app.delivery.send_to_providers.statsd_client")
     send_mock = mocker.patch("app.aws_ses_client.send_email", return_value="reference")
-    request_mock = mocker.patch(
-        "app.delivery.send_to_providers.urllib.request.Request",
-        return_value="request_mock",
+    mocker.patch("app.delivery.send_to_providers.Retry")
+
+    response_return_mock = MagicMock()
+    response_return_mock.status = 200
+    response_return_mock.data = "Hello there!"
+
+    response_mock = mocker.patch(
+        "app.delivery.send_to_providers.PoolManager.request",
+        return_value=response_return_mock,
     )
-    # See https://stackoverflow.com/a/34929900
-    cm = MagicMock()
-    cm.read.return_value = "request_content"
-    cm.__enter__.return_value = cm
-    cm.getcode = lambda: 200
-    urlopen_mock = mocker.patch("app.delivery.send_to_providers.urllib.request.urlopen")
-    urlopen_mock.return_value = cm
 
     send_to_providers.send_email_to_provider(db_notification)
 
     attachments = []
     if filename_attribute_present:
-        request_mock.assert_called_once_with("http://foo.bar/direct_file_url")
-        urlopen_mock.assert_called_once_with("request_mock")
+        response_mock.assert_called_with("GET", url="http://foo.bar/direct_file_url")
         attachments = [
             {
-                "data": "request_content",
                 "name": expected_filename,
+                "data": "Hello there!",
                 "mime_type": "application/pdf",
             }
         ]
@@ -1001,12 +1076,7 @@ def test_notification_with_bad_file_attachment_url(mocker, notify_db, notify_db_
 
     db_notification = save_notification(create_notification(template=template, personalisation=personalisation))
 
-    # See https://stackoverflow.com/a/34929900
-    cm = MagicMock()
-    cm.read.return_value = "request_content"
-    cm.__enter__.return_value = cm
-    urlopen_mock = mocker.patch("app.delivery.send_to_providers.urllib.request.urlopen")
-    urlopen_mock.return_value = cm
+    mocker.patch("app.delivery.send_to_providers.Retry")
 
     with pytest.raises(InvalidUrlException):
         send_to_providers.send_email_to_provider(db_notification)