From b787ca6e5b84bec9b4eef3f860fbb0aa5089f1e7 Mon Sep 17 00:00:00 2001
From: Alexey Bezhan <al@fueledbylemons.com>
Date: Mon, 3 Dec 2018 17:57:11 +0000
Subject: [PATCH] Add Service.get_days_of_retention helper method

Data retention lookup by type is only performed to get the number
of days, so we can update the service method to return the number
or the default directly.
---
 app/main/views/jobs.py          | 9 ++-------
 app/main/views/notifications.py | 5 +----
 app/models/service.py           | 6 +++---
 3 files changed, 6 insertions(+), 14 deletions(-)

diff --git a/app/main/views/jobs.py b/app/main/views/jobs.py
index ee8eeb102..36ba99da0 100644
--- a/app/main/views/jobs.py
+++ b/app/main/views/jobs.py
@@ -3,7 +3,6 @@
 from flask import (
     Response,
     abort,
-    current_app,
     jsonify,
     redirect,
     render_template,
@@ -226,9 +225,7 @@ def get_notifications(service_id, message_type, status_override=None):
     service_data_retention_days = None
 
     if message_type is not None:
-        service_data_retention_days = current_service.get_data_retention_by_type(
-            message_type
-        ).get('days_of_retention', current_app.config['ACTIVITY_STATS_LIMIT_DAYS'])
+        service_data_retention_days = current_service.get_days_of_retention(message_type)
 
     if request.path.endswith('csv') and current_user.has_permissions('view_activity'):
         return Response(
@@ -406,9 +403,7 @@ def get_job_partials(job, template):
             counts=_get_job_counts(job),
             status=filter_args['status']
         )
-    service_data_retention_days = current_service.get_data_retention_by_type(
-        template['template_type']
-    ).get('days_of_retention', current_app.config['ACTIVITY_STATS_LIMIT_DAYS'])
+    service_data_retention_days = current_service.get_days_of_retention(template['template_type'])
 
     return {
         'counts': counts,
diff --git a/app/main/views/notifications.py b/app/main/views/notifications.py
index 22251775a..6b08fc86c 100644
--- a/app/main/views/notifications.py
+++ b/app/main/views/notifications.py
@@ -8,7 +8,6 @@ from dateutil import parser
 from flask import (
     Response,
     abort,
-    current_app,
     jsonify,
     render_template,
     request,
@@ -197,9 +196,7 @@ def download_notifications_csv(service_id):
     filter_args = parse_filter_args(request.args)
     filter_args['status'] = set_status_filters(filter_args)
 
-    service_data_retention_days = current_service.get_data_retention_by_type(
-        filter_args.get('message_type')[0]
-    ).get('days_of_retention', current_app.config['ACTIVITY_STATS_LIMIT_DAYS'])
+    service_data_retention_days = current_service.get_days_of_retention(filter_args.get('message_type')[0])
     return Response(
         stream_with_context(
             generate_notifications_csv(
diff --git a/app/models/service.py b/app/models/service.py
index 4ed78ca50..aa86ac35e 100644
--- a/app/models/service.py
+++ b/app/models/service.py
@@ -1,4 +1,4 @@
-from flask import abort
+from flask import abort, current_app
 from notifications_utils.field import Field
 from werkzeug.utils import cached_property
 
@@ -273,11 +273,11 @@ class Service():
             None
         )
 
-    def get_data_retention_by_type(self, notification_type):
+    def get_days_of_retention(self, notification_type):
         return next(
             (dr for dr in self.data_retention if dr['notification_type'] == notification_type),
             {}
-        )
+        ).get('days_of_retention', current_app.config['ACTIVITY_STATS_LIMIT_DAYS'])
 
     @property
     def email_branding_id(self):