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

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/code_not_received.py            |  5 +++--
 tests/app/main/views/test_code_not_received.py | 18 ++++++++++++++----
 2 files changed, 17 insertions(+), 6 deletions(-)

diff --git a/app/main/views/code_not_received.py b/app/main/views/code_not_received.py
index 5b75325b1..a644e3f34 100644
--- a/app/main/views/code_not_received.py
+++ b/app/main/views/code_not_received.py
@@ -39,10 +39,11 @@ def check_and_resend_text_code():
 def check_and_resend_verification_code():
     user = User.from_email_address(session['user_details']['email'])
     user.send_verify_code()
+    redirect_url = request.args.get('next')
     if user.state == 'pending':
-        return redirect(url_for('main.verify'))
+        return redirect(url_for('main.verify', next=redirect_url))
     else:
-        return redirect(url_for('main.two_factor'))
+        return redirect(url_for('main.two_factor', next=redirect_url))
 
 
 @main.route('/email-not-received', methods=['GET'])
diff --git a/tests/app/main/views/test_code_not_received.py b/tests/app/main/views/test_code_not_received.py
index 9fd5452e3..143392987 100644
--- a/tests/app/main/views/test_code_not_received.py
+++ b/tests/app/main/views/test_code_not_received.py
@@ -119,27 +119,37 @@ def test_should_resend_verify_code_and_update_mobile_for_pending_user(
     )
 
 
+@pytest.mark.parametrize('redirect_url', [
+    None,
+    'blob',
+])
 def test_check_and_redirect_to_two_factor_if_user_active(
     client,
     api_user_active,
     mock_get_user_by_email,
     mock_send_verify_code,
+    redirect_url
 ):
     with client.session_transaction() as session:
         session['user_details'] = {
             'id': api_user_active['id'],
             'email': api_user_active['email_address']}
-    response = client.get(url_for('main.check_and_resend_verification_code'))
+    response = client.get(url_for('main.check_and_resend_verification_code', next=redirect_url))
     assert response.status_code == 302
-    assert response.location == url_for('main.two_factor', _external=True)
+    assert response.location == url_for('main.two_factor', _external=True, next=redirect_url)
 
 
+@pytest.mark.parametrize('redirect_url', [
+    None,
+    'blob',
+])
 def test_check_and_redirect_to_verify_if_user_pending(
     client,
     mocker,
     api_user_pending,
     mock_get_user_pending,
     mock_send_verify_code,
+    redirect_url
 ):
 
     mocker.patch('app.user_api_client.get_user_by_email', return_value=api_user_pending)
@@ -148,9 +158,9 @@ def test_check_and_redirect_to_verify_if_user_pending(
         session['user_details'] = {
             'id': api_user_pending['id'],
             'email': api_user_pending['email_address']}
-    response = client.get(url_for('main.check_and_resend_verification_code'))
+    response = client.get(url_for('main.check_and_resend_verification_code', next=redirect_url))
     assert response.status_code == 302
-    assert response.location == url_for('main.verify', _external=True)
+    assert response.location == url_for('main.verify', _external=True, next=redirect_url)
 
 
 @pytest.mark.parametrize('endpoint', [