From a6543c6e57f93bf1ff9f6c880ea0881cf011294a Mon Sep 17 00:00:00 2001
From: Pea Tyczynska <pea.tyczynska@digital.cabinet-office.gov.uk>
Date: Fri, 9 Oct 2020 11:41:06 +0100
Subject: [PATCH] Turn on redirects for sign_in

This is part of the work to make sure user is redirected
to the page they initially were meant to visit after
they sign in.
---
 app/main/views/sign_in.py            |  9 ++++----
 tests/app/main/views/test_sign_in.py | 31 +++++++++++++++++++++-------
 2 files changed, 28 insertions(+), 12 deletions(-)

diff --git a/app/main/views/sign_in.py b/app/main/views/sign_in.py
index d5b9cd1bb..7ce9225b4 100644
--- a/app/main/views/sign_in.py
+++ b/app/main/views/sign_in.py
@@ -25,6 +25,7 @@ def sign_in():
 
     form = LoginForm()
     password_reset_url = url_for('.forgot_password', next=request.args.get('next'))
+    redirect_url = request.args.get('next')
 
     if form.validate_on_submit():
 
@@ -33,7 +34,7 @@ def sign_in():
         )
 
         if user and user.state == 'pending':
-            return redirect(url_for('main.resend_email_verification'))
+            return redirect(url_for('main.resend_email_verification', next=redirect_url))
 
         if user and session.get('invited_user'):
             invited_user = InvitedUser.from_session()
@@ -45,9 +46,9 @@ def sign_in():
                 invited_user.accept_invite()
         if user and user.sign_in():
             if user.sms_auth:
-                return redirect(url_for('.two_factor', next=request.args.get('next')))
+                return redirect(url_for('.two_factor', next=redirect_url))
             if user.email_auth:
-                return redirect(url_for('.two_factor_email_sent'))
+                return redirect(url_for('.two_factor_email_sent', next=redirect_url))
 
         # Vague error message for login in case of user not known, locked, inactive or password not verified
         flash(Markup(
@@ -61,7 +62,7 @@ def sign_in():
     return render_template(
         'views/signin.html',
         form=form,
-        again=bool(request.args.get('next')),
+        again=bool(redirect_url),
         other_device=other_device,
         password_reset_url=password_reset_url
     )
diff --git a/tests/app/main/views/test_sign_in.py b/tests/app/main/views/test_sign_in.py
index 8e0208f20..e22d929db 100644
--- a/tests/app/main/views/test_sign_in.py
+++ b/tests/app/main/views/test_sign_in.py
@@ -102,6 +102,10 @@ def test_logged_in_user_redirects_to_account(
     )
 
 
+@pytest.mark.parametrize('redirect_url', [
+    None,
+    'blob',
+])
 @pytest.mark.parametrize('email_address, password', [
     ('valid@example.gov.uk', 'val1dPassw0rd!'),
     (' valid@example.gov.uk  ', '  val1dPassw0rd!  '),
@@ -115,34 +119,40 @@ def test_process_sms_auth_sign_in_return_2fa_template(
     mock_verify_password,
     email_address,
     password,
+    redirect_url
 ):
     response = client.post(
-        url_for('main.sign_in'), data={
+        url_for('main.sign_in', next=redirect_url), data={
             'email_address': email_address,
             'password': password})
     assert response.status_code == 302
-    assert response.location == url_for('.two_factor', _external=True)
+    assert response.location == url_for('.two_factor', next=redirect_url, _external=True)
     mock_verify_password.assert_called_with(api_user_active['id'], password)
     mock_get_user_by_email.assert_called_with('valid@example.gov.uk')
 
 
+@pytest.mark.parametrize('redirect_url', [
+    None,
+    'blob',
+])
 def test_process_email_auth_sign_in_return_2fa_template(
     client,
     api_user_active_email_auth,
     mock_send_verify_code,
     mock_verify_password,
-    mocker
+    mocker,
+    redirect_url
 ):
     mocker.patch('app.user_api_client.get_user', return_value=api_user_active_email_auth)
     mocker.patch('app.user_api_client.get_user_by_email', return_value=api_user_active_email_auth)
 
     response = client.post(
-        url_for('main.sign_in'), data={
+        url_for('main.sign_in', next=redirect_url), data={
             'email_address': 'valid@example.gov.uk',
             'password': 'val1dPassw0rd!'})
     assert response.status_code == 302
-    assert response.location == url_for('.two_factor_email_sent', _external=True)
-    mock_send_verify_code.assert_called_with(api_user_active_email_auth['id'], 'email', None, None)
+    assert response.location == url_for('.two_factor_email_sent', _external=True, next=redirect_url)
+    mock_send_verify_code.assert_called_with(api_user_active_email_auth['id'], 'email', None, redirect_url)
     mock_verify_password.assert_called_with(api_user_active_email_auth['id'], 'val1dPassw0rd!')
 
 
@@ -185,16 +195,21 @@ def test_should_return_redirect_when_user_is_pending(
     assert response.status_code == 200
 
 
+@pytest.mark.parametrize('redirect_url', [
+    None,
+    'blob',
+])
 def test_should_attempt_redirect_when_user_is_pending(
     client,
     mock_get_user_by_email_pending,
     mock_verify_password,
+    redirect_url
 ):
     response = client.post(
-        url_for('main.sign_in'), data={
+        url_for('main.sign_in', next=redirect_url), data={
             'email_address': 'pending_user@example.gov.uk',
             'password': 'val1dPassw0rd!'})
-    assert response.location == url_for('main.resend_email_verification', _external=True)
+    assert response.location == url_for('main.resend_email_verification', _external=True, next=redirect_url)
     assert response.status_code == 302