From 2676ee9bcf66280d361e82434710fc40d0166b72 Mon Sep 17 00:00:00 2001
From: Imdad Ahad <imdad.ahad@digital.cabinet-office.gov.uk>
Date: Fri, 28 Oct 2016 10:47:32 +0100
Subject: [PATCH] Add additional tests and refactor

---
 app/main/views/add_service.py           | 58 +++++++++++++++----------
 app/main/views/choose_service.py        |  8 ++--
 app/templates/views/choose-service.html | 12 ++---
 3 files changed, 43 insertions(+), 35 deletions(-)

diff --git a/app/main/views/add_service.py b/app/main/views/add_service.py
index 7cc588e54..3f8a5e064 100644
--- a/app/main/views/add_service.py
+++ b/app/main/views/add_service.py
@@ -1,5 +1,3 @@
-import re
-
 from flask import (
     render_template,
     redirect,
@@ -27,34 +25,49 @@ from app import (
 
 from app.utils import (
     email_safe,
-    user_in_whitelist
+    is_gov_user
 )
 
 
+def _add_invited_user_to_service(invited_user):
+    invitation = InvitedUser(**invited_user)
+    # if invited user add to service and redirect to dashboard
+    user = user_api_client.get_user(session['user_id'])
+    service_id = invited_user['service']
+    user_api_client.add_user_to_service(service_id, user.id, invitation.permissions)
+    invite_api_client.accept_invite(service_id, invitation.id)
+    return service_id
+
+
+def _create_service(service_name, email_from):
+    service_id = service_api_client.create_service(service_name=service_name,
+                                                   active=False,
+                                                   message_limit=current_app.config['DEFAULT_SERVICE_LIMIT'],
+                                                   restricted=True,
+                                                   user_id=session['user_id'],
+                                                   email_from=email_from)
+    session['service_id'] = service_id
+    return service_id
+
+
 @main.route("/add-service", methods=['GET', 'POST'])
 @login_required
 def add_service():
     invited_user = session.get('invited_user')
     if invited_user:
-        invitation = InvitedUser(**invited_user)
-        # if invited user add to service and redirect to dashboard
-        user = user_api_client.get_user(session['user_id'])
-        service_id = invited_user['service']
-        user_api_client.add_user_to_service(service_id, user.id, invitation.permissions)
-        invite_api_client.accept_invite(service_id, invitation.id)
+        service_id = _add_invited_user_to_service(invited_user)
         return redirect(url_for('main.service_dashboard', service_id=service_id))
 
+    if not is_gov_user(current_user.email_address):
+        abort(403)
+
     form = AddServiceForm(service_api_client.find_all_service_email_from)
     heading = 'Which service do you want to set up notifications for?'
+
     if form.validate_on_submit():
         email_from = email_safe(form.name.data)
-        service_id = service_api_client.create_service(service_name=form.name.data,
-                                                       active=False,
-                                                       message_limit=current_app.config['DEFAULT_SERVICE_LIMIT'],
-                                                       restricted=True,
-                                                       user_id=session['user_id'],
-                                                       email_from=email_from)
-        session['service_id'] = service_id
+        service_name = form.name.data
+        service_id = _create_service(service_name, email_from)
 
         if (len(service_api_client.get_services({'user_id': session['user_id']}).get('data', [])) > 1):
             return redirect(url_for('main.service_dashboard', service_id=service_id))
@@ -73,11 +86,8 @@ def add_service():
             help=1
         ))
     else:
-        if not user_in_whitelist(current_user.email_address):
-            abort(403)
-        else:
-            return render_template(
-                'views/add-service.html',
-                form=form,
-                heading=heading
-            )
+        return render_template(
+            'views/add-service.html',
+            form=form,
+            heading=heading
+        )
diff --git a/app/main/views/choose_service.py b/app/main/views/choose_service.py
index 48bf81e39..d5cd2e6b9 100644
--- a/app/main/views/choose_service.py
+++ b/app/main/views/choose_service.py
@@ -1,11 +1,9 @@
-import re
-
-from flask import (current_app, render_template, redirect, url_for, session)
+from flask import (render_template, redirect, url_for, session)
 from flask_login import login_required, current_user
 from app.main import main
 from app import service_api_client
 from app.notify_client.service_api_client import ServicesBrowsableItem
-from app.utils import user_in_whitelist
+from app.utils import is_gov_user
 
 
 @main.route("/services")
@@ -15,7 +13,7 @@ def choose_service():
         'views/choose-service.html',
         services=[ServicesBrowsableItem(x) for x in
                   service_api_client.get_services({'user_id': current_user.id})['data']],
-        can_add_service=user_in_whitelist(current_user.email_address)
+        can_add_service=is_gov_user(current_user.email_address)
     )
 
 
diff --git a/app/templates/views/choose-service.html b/app/templates/views/choose-service.html
index f55bda167..fbc55ba4d 100644
--- a/app/templates/views/choose-service.html
+++ b/app/templates/views/choose-service.html
@@ -13,12 +13,12 @@
 
   {{ browse_list(services) }}
   {% if can_add_service %}
-  {{ browse_list([
-    {
-      'title': 'Add a new service…',
-      'link': url_for('.add_service')
-    },
-  ]) }}
+    {{ browse_list([
+      {
+        'title': 'Add a new service…',
+        'link': url_for('.add_service')
+      },
+    ]) }}
   {% endif %}